| 25 | |
| 26 | |
| 27 | class EmbedxNode: |
| 28 | |
| 29 | def __init__(self, json_node): |
| 30 | self.json_node = json_node |
| 31 | |
| 32 | @property |
| 33 | def node_id(self): |
| 34 | assert "id" in self.json_node |
| 35 | return self.json_node["id"] |
| 36 | |
| 37 | @property |
| 38 | def label(self): |
| 39 | assert "label" in self.json_node |
| 40 | return self.json_node["label"] |
| 41 | |
| 42 | @property |
| 43 | def stage(self): |
| 44 | assert "stage" in self.json_node |
| 45 | return self.json_node["stage"] |
| 46 | |
| 47 | @property |
| 48 | def feature(self): |
| 49 | assert "feature" in self.json_node or "feat" in self.json_node |
| 50 | if "feature" in self.json_node: |
| 51 | return self.json_node["feature"] |
| 52 | return self.json_node["feat"] |
| 53 | |
| 54 | @property |
| 55 | def valid_feature(self): |
| 56 | return self.feature |
| 57 | |
| 58 | def has_feature(self): |
| 59 | return "feature" in self.json_node or "feat" in self.json_node |
| 60 | |
| 61 | @classmethod |
| 62 | def get_node_feature_str(cls, node): |
| 63 | if not node.has_feature(): |
| 64 | return "" |
| 65 | |
| 66 | valid_feature_str = " ".join( |
| 67 | [f"{i}:{feat}" for i, feat in enumerate(node.valid_feature) if feat != 0]) |
| 68 | return f"{node.node_id} {valid_feature_str}" |
| 69 | |
| 70 | @classmethod |
| 71 | def get_node_label_str(cls, node): |
| 72 | node_str = str(node.node_id) |
| 73 | label = node.label |
| 74 | if isinstance(label, int): |
| 75 | node_str += " " + str(label) |
| 76 | elif isinstance(label, list): |
| 77 | node_str += " " + " ".join([str(int(x)) for x in label]) |
| 78 | elif isinstance(label, np.ndarray): |
| 79 | label_nnz = np.nonzero(label)[0] |
| 80 | if np.isscalar(label_nnz): |
| 81 | node_str += " " + str(np.asscalar(label_nnz.a)) |
| 82 | else: |
| 83 | node_str += " " + " ".join([str(int(x)) for x in label]) |
| 84 | return node_str |