Add generator_id parameter to bayesdb_variable_numbers.

Require callers to specify whether they want to include the latent
variables of a specific generator or not.

Author: riastradh-probcomp

src/bqlfn.py

@@ -356,7 +356,8 @@ def bql_row_similarity(bdb, population_id, generator_id, rowid, target_rowid,
     if target_rowid is None:
         raise BQLError(bdb, 'No such target row for SIMILARITY')
     if len(colnos) == 0:
-        colnos = core.bayesdb_variable_numbers(bdb, population_id)
+        colnos = core.bayesdb_variable_numbers(bdb, population_id,
+            generator_id)
     def generator_similarity(generator_id):
         metamodel = core.bayesdb_generator_metamodel(bdb, generator_id)
         return metamodel.row_similarity(bdb, generator_id, None, rowid,

src/compiler.py

@@ -1194,7 +1194,8 @@ def compile_column_lists(bdb, population_id, generator_id, column_lists,
         else:
             out.write(', ')
         if isinstance(collist, ast.ColListAll):
-            colnos = core.bayesdb_variable_numbers(bdb, population_id)
+            colnos = core.bayesdb_variable_numbers(bdb, population_id,
+                generator_id)
             out.write(', '.join(str(colno) for colno in colnos))
         elif isinstance(collist, ast.ColListLit):
             unknown = set()

src/core.py

@@ -244,20 +244,27 @@ def bayesdb_variable_number(bdb, population_id, generator_id, name):
         ''', (population_id, generator_id, name))
     return cursor_value(cursor)
 
-def bayesdb_variable_names(bdb, population_id):
+def bayesdb_variable_names(bdb, population_id, generator_id):
     """Return a list of the names of columns modelled in `population_id`."""
-    colnos = bayesdb_variable_numbers(bdb, population_id)
+    colnos = bayesdb_variable_numbers(bdb, population_id, generator_id)
     return [bayesdb_variable_name(bdb, population_id, colno)
         for colno in colnos]
 
-def bayesdb_variable_numbers(bdb, population_id):
+def bayesdb_variable_numbers(bdb, population_id, generator_id):
     """Return a list of the numbers of columns modelled in `population_id`."""
-    sql = '''
-        SELECT colno FROM bayesdb_variable
-            WHERE population_id = ? AND generator_id IS NULL
-            ORDER BY colno ASC
-    '''
-    return [row[0] for row in bdb.sql_execute(sql, (population_id,))]
+    if generator_id is None:
+        cursor = bdb.sql_execute('''
+            SELECT colno FROM bayesdb_variable
+                WHERE population_id = ? AND generator_id IS NULL
+                ORDER BY colno ASC
+        ''', (population_id,))
+    else:
+        cursor = bdb.sql_execute('''
+            SELECT colno FROM bayesdb_variable
+                WHERE population_id = ? AND generator_id = ?
+                ORDER BY colno ASC
+        ''', (population_id, generator_id))
+    return [colno for (colno,) in cursor]
 
 def bayesdb_variable_name(bdb, population_id, colno):
     """Return the name a population variable."""

src/metamodels/cgpm_metamodel.py

@@ -237,7 +237,8 @@ def analyze_models(self, bdb, generator_id, modelnos=None, iterations=1,
         def retrieve_analyze_variables(ast):
             # Transition all variables by default.
             if len(ast) == 0:
-                variables = core.bayesdb_variable_names(bdb, population_id)
+                variables = core.bayesdb_variable_names(bdb, population_id,
+                    generator_id)
             # Exactly 1 clause supported.
             elif len(ast) == 1:
                 clause = ast[0]
@@ -248,7 +249,8 @@ def retrieve_analyze_variables(ast):
                 elif isinstance(clause, cgpm_analyze.parse.Skip):
                     variables = filter(
                         lambda v: v not in clause.vars,
-                        core.bayesdb_variable_names(bdb, population_id))
+                        core.bayesdb_variable_names(bdb, population_id,
+                            generator_id))
                 # Unknown/impossible clause.
                 else:
                     raise ValueError('Unknown clause in ANALYZE: %s.' % ast)
@@ -672,7 +674,7 @@ def _retrieve_stattype_dist_params(var):
 
             # Get the column number.
             colno = core.bayesdb_variable_number(bdb, population_id, None, var)
-            assert not colno < 0
+            assert 0 <= colno
 
             # Add it to the list and mark it modelled by default.
             stattype = core.bayesdb_variable_stattype(
@@ -785,12 +787,11 @@ def _retrieve_stattype_dist_params(var):
     # Use the default distribution for any variables that remain to be
     # modelled, excluding any that are latent or that have statistical
     # types we don't know about.
-    for var in core.bayesdb_variable_names(bdb, population_id):
+    for var in core.bayesdb_variable_names(bdb, population_id, None):
         if var in modelled:
             continue
         colno = core.bayesdb_variable_number(bdb, population_id, None, var)
-        if colno < 0:
-            continue
+        assert 0 <= colno
         stattype, dist, params = _retrieve_stattype_dist_params(var)
         if stattype not in _DEFAULT_DIST:
             assert False    # XXX Why would you be here, anyway?