def _sample_nodes(model: GraphSamplingModel, graph, node_feature_mapping, generator=None): labels_idx = model.sample_node_label(graph, generator) if len(labels_idx) != 1: raise NotImplementedError('Batch sampling is not implemented') label_idx = int(labels_idx[0]) label = dataset.NODE_IDX_MAP_REVERSE[label_idx] target_type = target.TargetType.from_label(label) if target_type not in model.feature_dimensions: return [NodeOp(label)] if target_type in model.feature_dimensions: fd = model.feature_dimensions[target_type] numerical_features = labels_idx.new_zeros((1, len(fd)), dtype=int) for feat_idx in range(len(fd)): sample_features = model.sample_entity_features( graph, target_type, numerical_features, generator) numerical_features[0, feat_idx] = sample_features[feat_idx] numerical_features = numerical_features[0].cpu().tolist() parameters = node_feature_mapping.features_from_index( numerical_features, target_type) else: parameters = {} return [NodeOp(label, parameters)]
def _add_node(self, label, parameters=None): if parameters is None: parameters = {} self.seq.append(NodeOp(label, parameters)) self.num_nodes += 1 if not isinstance(label, datalib.SubnodeType): self.last_entity_node = self.num_nodes - 1 return # Add subedge self.add_op( EdgeOp(datalib.ConstraintType.Subnode, (self.num_nodes - 1, self.last_entity_node)))
def test_get_sequence_dof(): seq = [ NodeOp(label=EntityType.External), NodeOp(label=EntityType.Line), NodeOp(label=SubnodeType.SN_Start), EdgeOp(label=ConstraintType.Subnode, references=(2, 1)), NodeOp(label=SubnodeType.SN_End), EdgeOp(label=ConstraintType.Subnode, references=(3, 1)), NodeOp(label=EntityType.Line), EdgeOp(label=ConstraintType.Parallel, references=(4, 1)), EdgeOp(label=ConstraintType.Horizontal, references=(4,)), EdgeOp(label=ConstraintType.Distance, references=(4, 1)), NodeOp(label=SubnodeType.SN_Start), EdgeOp(label=ConstraintType.Subnode, references=(5, 4)), NodeOp(label=SubnodeType.SN_End), EdgeOp(label=ConstraintType.Subnode, references=(6, 4)), NodeOp(label=EntityType.Stop)] dof_remaining = np.sum(get_sequence_dof(seq)) assert dof_remaining == 5
def generate_sample(model: GraphSamplingModel, max_iters, node_feature_mapping, edge_feature_mapping, generator=None, device=None): builder = _SeqBuilder() state = 'add_node' subnodes_to_add = None while len(builder.seq) < max_iters: graph = dataset.graph_info_from_sequence(builder.seq, node_feature_mapping, edge_feature_mapping) graph = training.load_cuda_async(graph, device) if state == 'add_node': node_op, = _sample_nodes(model, graph, node_feature_mapping, generator) builder.add_op(node_op) if node_op.label == EntityType.Stop: break subnodes_to_add = list(_get_subnodes_for_entity(node_op)) state = 'add_edge' elif state == 'add_edge': edge_op, = _sample_edges(model, graph, edge_feature_mapping, generator) if edge_op is not None: assert max(edge_op.references) + 1 == builder.num_nodes builder.add_op(edge_op) continue if subnodes_to_add: subnode_op = NodeOp(subnodes_to_add.pop()) builder.add_op(subnode_op) state = 'add_edge' else: state = 'add_node' else: assert False return builder.seq
def __init__(self): self.seq = [NodeOp(datalib.EntityType.External)] self.last_entity_node = None self.num_nodes = 1