Handle distribution parameters when parsing hirm schema from python

cxx/assets/animals.unary.hirm

@@ -0,0 +1,23 @@
+0 quadrapedal strainteeth oldworld ocean coastal hairless flippers black flys skimmer buckteeth plankton walks plains fish blue swims arctic forager ground patches furry fields smelly water white
+5 paws tunnels muscle longleg group newworld brown nocturnal agility hibernate active bulbous fast lean inactive small stripes chewteeth strong tail toughskin weak big slow
+7 pads vegetation meatteeth hooves longneck hunter smart claws insects meat hops timid horns grazer orange gray scavenger spots stalker yellow fierce tusks
+8 mountains jungle nestspot hands domestic bush bipedal red desert solitary forest cave tree
+
+irm=0
+animal 0 antelope bat bobcat buffalo chihuahua chimpanzee collie cow dalmatian deer elephant fox germanshepherd giantpanda giraffe gorilla grizzlybear hamster horse leopard lion mole moose mouse ox persiancat pig rabbit raccoon rat rhinoceros sheep siamesecat skunk spidermonkey squirrel tiger weasel wolf zebra
+animal 1 bluewhale dolphin humpbackwhale killerwhale seal walrus
+animal 2 beaver hippopotamus otter polarbear
+
+irm=5
+animal 4 antelope chimpanzee dalmatian deer dolphin germanshepherd giraffe gorilla horse killerwhale leopard lion seal spidermonkey tiger wolf zebra
+animal 5 bluewhale buffalo cow elephant giantpanda grizzlybear hippopotamus humpbackwhale moose ox pig polarbear rhinoceros sheep walrus
+animal 6 bat beaver bobcat chihuahua collie fox hamster mole mouse otter persiancat rabbit raccoon rat siamesecat skunk squirrel weasel
+
+irm=7
+animal 0 antelope beaver bluewhale buffalo chimpanzee cow deer dolphin elephant giantpanda giraffe hamster hippopotamus horse humpbackwhale mole moose mouse ox pig rabbit rhinoceros seal sheep skunk spidermonkey squirrel walrus zebra
+animal 1 bat bobcat chihuahua collie dalmatian fox germanshepherd gorilla grizzlybear killerwhale leopard lion otter persiancat polarbear raccoon rat siamesecat tiger weasel wolf
+
+irm=8
+animal 0 antelope beaver bluewhale buffalo chihuahua collie cow dalmatian deer dolphin elephant fox germanshepherd giraffe hamster hippopotamus horse humpbackwhale killerwhale mole moose mouse otter ox persiancat pig polarbear rabbit rat rhinoceros seal sheep siamesecat skunk tiger walrus weasel zebra
+animal 9 bat chimpanzee giantpanda gorilla spidermonkey squirrel
+animal 12 bobcat grizzlybear leopard lion raccoon wolf

cxx/assets/pclean_rents_clean_ternary.schema

@@ -1,3 +1,3 @@
-has_type ~ stringcat[strings="1br 2br 3br 4br studio"](record County State)
+has_type ~ stringcat[strings="1br 2br 3br 4br studio"](record, County, State)
 has_rent ~ normal(record, County, State)
 county_name ~ bigram(County, State)

py/hirm_io.py

@@ -22,8 +22,7 @@ def load_schema(path):
   """Load the schema from path, and return it as a dict of Relations."""
   relations = {}
   comment_line_re = re.compile(r'\s*#.*')
-  # TODO(thomaswc): Handle distribution parameters in brackets
-  line_re = re.compile(r'\s*(\w+)\s*~\s*(\w+)\s*\(\s*([^\)]+)\s*\)\s*(#.*)?')
+  line_re = re.compile(r'\s*(\w+)\s*~\s*(\w+)\s*(\[[^\]]*\])?\(\s*([^\)]+)\s*\)\s*(#.*)?')
   with open(path, 'r') as f:
     for line in f:
       if comment_line_re.match(line):
@@ -33,10 +32,18 @@ def load_schema(path):
         print(f"Could not parse schema line\n{line}\n")
         sys.exit(1)
       name = m.group(1)
-      domains = re.split(r'\s*,\s*', m.group(3))
+      parameters = {}
+      if m.group(3):
+        param_string = m.group(3)[1:-1]
+        ps = re.split(r'\s*,\s*', param_string)
+        for p in ps:
+          kv = re.split(r'\s*=\s*', p)
+          assert len(kv) == 2
+          parameters[kv[0]] = kv[1]
+      domains = re.split(r'\s*,\s*', m.group(4))
       relations[name] = Relation(name=name,
                                  distribution=m.group(2),
-                                 parameters={},
+                                 parameters=parameters,
                                  domains=domains)
   return relations
 

py/hirm_io_test.py

@@ -22,6 +22,16 @@ def test_load_schema_with_comments(self):
     self.assertEqual(relations["big"].distribution, "bernoulli")
     self.assertEqual(relations["big"].domains, ["animal"])
 
+  def test_load_schema_with_distribution_parameters(self):
+    relations = hirm_io.load_schema("../cxx/assets/pclean_rents_clean_ternary.schema")
+    self.assertEqual(len(relations), 3)
+    self.assertEqual(relations["has_type"].name, "has_type")
+    self.assertEqual(relations["has_type"].distribution, "stringcat")
+    self.assertEqual(len(relations["has_type"].parameters), 1)
+    self.assertEqual(relations["has_type"].parameters["strings"],
+                     '"1br 2br 3br 4br studio"')
+    self.assertEqual(relations["has_type"].domains, ["record", "County", "State"])
+
   def test_load_observations(self):
     observations = hirm_io.load_observations("../cxx/assets/animals.unary.obs")
     self.assertEqual(len(observations), 4250)
@@ -31,15 +41,16 @@ def test_load_observations(self):
 
   def test_load_clusters(self):
     clusters = hirm_io.load_clusters("../cxx/assets/animals.unary.hirm")
-    self.assertEqual(len(clusters), 3)
-    self.assertEqual(clusters[0].cluster_id, "1")
-    self.assertEqual(clusters[1].cluster_id, "3")
-    self.assertEqual(clusters[2].cluster_id, "6")
-    self.assertEqual(clusters[0].relations[0], "paws")
+    self.assertEqual(len(clusters), 4)
+    self.assertEqual(clusters[0].cluster_id, "0")
+    self.assertEqual(clusters[1].cluster_id, "5")
+    self.assertEqual(clusters[2].cluster_id, "7")
+    self.assertEqual(clusters[3].cluster_id, "8")
+    self.assertEqual(clusters[0].relations[0], "quadrapedal")
     self.assertEqual(len(clusters[0].domain_clusters), 3)
     self.assertEqual(clusters[0].domain_clusters[0].cluster_id, "0")
     self.assertEqual(clusters[0].domain_clusters[0].domain, "animal")
-    self.assertEqual(clusters[0].domain_clusters[0].entities[0], "beaver")
+    self.assertEqual(clusters[0].domain_clusters[0].entities[0], "antelope")
 
 
 if __name__ == "__main__":