Skip to content

Commit

Permalink
feat: Allow adding/appending to groups when adding nodes or edges (#203)
Browse files Browse the repository at this point in the history
  • Loading branch information
@MarIniOnz
MarIniOnz authored Sep 17, 2024
1 parent d9e1d3c commit 3630bff
Show file tree
Hide file tree
Showing 4 changed files with 405 additions and 77 deletions.
47 changes: 4 additions & 43 deletions medmodels/medrecord/builder.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -178,46 +178,13 @@ def build(self) -> mm.MedRecord:
continue

group = node[1]

if not medrecord.contains_group(group):
medrecord.add_group(group)

node = node[0]

if is_node_tuple(node):
medrecord.add_node(*node)
medrecord.add_node_to_group(group, node[0])
continue

if is_node_tuple_list(node):
medrecord.add_nodes(node)
medrecord.add_node_to_group(group, [node[0] for node in node])
continue

if is_pandas_node_dataframe_input(node):
medrecord.add_nodes(node)
medrecord.add_node_to_group(group, node[0][node[1]].tolist())
continue

if is_polars_node_dataframe_input(node):
medrecord.add_nodes(node)
medrecord.add_node_to_group(group, node[0][node[1]].to_list())
continue

if is_pandas_node_dataframe_input_list(node):
medrecord.add_nodes(node)
medrecord.add_node_to_group(
group,
[nodes for node in node for nodes in node[0][node[1]].tolist()],
)
medrecord.add_node(*node, group)
continue

if is_polars_node_dataframe_input_list(node):
medrecord.add_nodes(node)
medrecord.add_node_to_group(
group,
[nodes for node in node for nodes in node[0][node[1]].to_list()],
)
medrecord.add_nodes(node, group)

for edge in self._edges:
if is_edge_tuple(edge):
Expand All @@ -235,15 +202,10 @@ def build(self) -> mm.MedRecord:
continue

group = edge[1]

if not medrecord.contains_group(group):
medrecord.add_group(group)

edge = edge[0]

if is_edge_tuple(edge):
edge_index = medrecord.add_edge(*edge)
medrecord.add_edge_to_group(group, edge_index)
medrecord.add_edge(*edge, group)
continue

if (
Expand All @@ -253,8 +215,7 @@ def build(self) -> mm.MedRecord:
or is_polars_edge_dataframe_input(edge)
or is_polars_edge_dataframe_input_list(edge)
):
edge_indices = medrecord.add_edges(edge)
medrecord.add_edge_to_group(group, edge_indices)
medrecord.add_edges(edge, group)

for group in self._groups:
if medrecord.contains_group(group):
Expand Down
160 changes: 129 additions & 31 deletions medmodels/medrecord/medrecord.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -533,17 +533,31 @@ def edges_connecting(
(target_node if isinstance(target_node, list) else [target_node]),
)

def add_node(self, node: NodeIndex, attributes: AttributesInput) -> None:
"""Adds a node with specified attributes to the MedRecord instance.
def add_node(
self,
node: NodeIndex,
attributes: AttributesInput,
group: Optional[Group] = None,
) -> None:
"""Adds a node with specified attributes to the MedRecord instance. Optionally adds the node to a group.
Args:
node (NodeIndex): The index of the node to add.
attributes (Attributes): A dictionary of the node's attributes.
group (Optional[Group]): The name of the group to add the node to, optional.
Returns:
None
"""
return self._medrecord.add_node(node, attributes)
self._medrecord.add_node(node, attributes)

if group is None:
return

if not self.contains_group(group):
self.add_group(group)

self.add_node_to_group(group, node)

@overload
def remove_node(self, node: NodeIndex) -> Attributes: ...
Expand Down Expand Up @@ -591,89 +605,140 @@ def add_nodes(
PolarsNodeDataFrameInput,
List[PolarsNodeDataFrameInput],
],
group: Optional[Group] = None,
) -> None:
"""Adds multiple nodes to the MedRecord from different data formats.
"""Adds multiple nodes to the MedRecord from different data formats and optionally assigns them to a group.
Accepts a list of tuples, DataFrame(s), or PolarsNodeDataFrameInput(s) to add
nodes. If a DataFrame or list of DataFrames is used, the add_nodes_pandas method
is called. If PolarsNodeDataFrameInput(s) are provided, each tuple must include
a DataFrame and the index column.
a DataFrame and the index column. If a group is specified, the nodes are added
to the group.
Args:
nodes (Union[Sequence[NodeTuple], PandasNodeDataFrameInput, List[PandasNodeDataFrameInput], PolarsNodeDataFrameInput, List[PolarsNodeDataFrameInput]]):
Data representing nodes in various formats.
group (Optional[Group]): The name of the group to add the nodes to. If not
specified, the nodes are added to the MedRecord without a group.
Returns:
None
"""
if is_pandas_node_dataframe_input(nodes) or is_pandas_node_dataframe_input_list(
nodes
):
return self.add_nodes_pandas(nodes)
self.add_nodes_pandas(nodes, group)
elif is_polars_node_dataframe_input(
nodes
) or is_polars_node_dataframe_input_list(nodes):
return self.add_nodes_polars(nodes)
self.add_nodes_polars(nodes, group)
else:
return self._medrecord.add_nodes(nodes)
self._medrecord.add_nodes(nodes)

if group is None:
return

if not self.contains_group(group):
self.add_group(group)

self.add_node_to_group(group, [node[0] for node in nodes])

def add_nodes_pandas(
self, nodes: Union[PandasNodeDataFrameInput, List[PandasNodeDataFrameInput]]
self,
nodes: Union[PandasNodeDataFrameInput, List[PandasNodeDataFrameInput]],
group: Optional[Group] = None,
) -> None:
"""Adds nodes to the MedRecord instance from one or more Pandas DataFrames.
"""Adds nodes to the MedRecord instance from one or more Pandas DataFrames. Optionally assigns them to a group.
This method accepts either a single tuple or a list of tuples, where each tuple
consists of a Pandas DataFrame and an index column string.
consists of a Pandas DataFrame and an index column string. If a group is
specified, the nodes are added to the group.
Args:
nodes (Union[PandasNodeDataFrameInput, List[PandasNodeDataFrameInput]]):
A tuple or list of tuples, each with a DataFrame and index column.
group (Optional[Group]): The name of the group to add the nodes to. If not
specified, the nodes are added to the MedRecord without a group.
Returns:
None
"""
return self.add_nodes_polars(
self.add_nodes_polars(
[process_nodes_dataframe(nodes_df) for nodes_df in nodes]
if isinstance(nodes, list)
else [process_nodes_dataframe(nodes)]
else [process_nodes_dataframe(nodes)],
group,
)

def add_nodes_polars(
self, nodes: Union[PolarsNodeDataFrameInput, List[PolarsNodeDataFrameInput]]
self,
nodes: Union[PolarsNodeDataFrameInput, List[PolarsNodeDataFrameInput]],
group: Optional[Group] = None,
) -> None:
"""Adds nodes to the MedRecord instance from one or more Polars DataFrames.
"""Adds nodes to the MedRecord instance from one or more Polars DataFrames. Optionally assigns them to a group.
This method accepts either a single tuple or a list of tuples, where each tuple
consists of a Polars DataFrame and an index column string.
consists of a Polars DataFrame and an index column string. If a group is
specified, the nodes are added to the group.
Args:
nodes (Union[PolarsNodeDataFrameInput, List[PolarsNodeDataFrameInput]]):
A tuple or list of tuples, each with a DataFrame and index column.
group (Optional[Group]): The name of the group to add the nodes to. If not
specified, the nodes are added to the MedRecord without a group.
Returns:
None
"""
return self._medrecord.add_nodes_dataframes(
self._medrecord.add_nodes_dataframes(
nodes if isinstance(nodes, list) else [nodes]
)

if group is None:
return

if not self.contains_group(group):
self.add_group(group)

if isinstance(nodes, list):
node_indices = [
nodes for node in nodes for nodes in node[0][node[1]].to_list()
]
else:
node_indices = nodes[0][nodes[1]].to_list()

self.add_node_to_group(group, node_indices)

def add_edge(
self,
source_node: NodeIndex,
target_node: NodeIndex,
attributes: AttributesInput,
group: Optional[Group] = None,
) -> EdgeIndex:
"""Adds an edge between two specified nodes with given attributes.
"""Adds an edge between two specified nodes with given attributes. Optionally assigns the edge to a group.
Args:
source_node (NodeIndex): Index of the source node.
target_node (NodeIndex): Index of the target node.
attributes (AttributesInput): Dictionary or mapping of edge attributes.
group (Optional[Group]): The name of the group to add the edge to. If not
specified, the edge is added to the MedRecord without a group.
Returns:
EdgeIndex: The index of the added edge.
"""
return self._medrecord.add_edge(source_node, target_node, attributes)
edge_index = self._medrecord.add_edge(source_node, target_node, attributes)

if group is None:
return edge_index

if not self.contains_group(group):
self.add_group(group)

self.add_edge_to_group(group, edge_index)

return edge_index

@overload
def remove_edge(self, edge: EdgeIndex) -> Attributes: ...
Expand Down Expand Up @@ -721,78 +786,111 @@ def add_edges(
PolarsEdgeDataFrameInput,
List[PolarsEdgeDataFrameInput],
],
group: Optional[Group] = None,
) -> List[EdgeIndex]:
"""Adds edges to the MedRecord instance from various data formats.
"""Adds edges to the MedRecord instance from various data formats. Optionally assigns them to a group.
Accepts lists of tuples, DataFrame(s), or EdgeDataFrameInput(s) to add edges.
Each tuple must have indices for source and target nodes and a dictionary of
attributes. If a DataFrame or list of DataFrames is used,
the add_edges_dataframe method is invoked.
the add_edges_dataframe method is invoked. If PolarsEdgeDataFrameInput(s) are
provided, each tuple must include a DataFrame and index columns for source and
target nodes. If a group is specified, the edges are added to the group.
Args:
edges (Union[Sequence[EdgeTuple], PandasEdgeDataFrameInput, List[PolarsEdgeDataFrameInput]]):
List[PandasEdgeDataFrameInput], PolarsEdgeDataFrameInput,
Data representing edges in several formats.
group (Optional[Group]): The name of the group to add the edges to. If not
specified, the edges are added to the MedRecord without a group.
Returns:
List[EdgeIndex]: A list of edge indices that were added.
"""
if is_pandas_edge_dataframe_input(edges) or is_pandas_edge_dataframe_input_list(
edges
):
return self.add_edges_pandas(edges)
return self.add_edges_pandas(edges, group)
elif is_polars_edge_dataframe_input(
edges
) or is_polars_edge_dataframe_input_list(edges):
return self.add_edges_polars(edges)
return self.add_edges_polars(edges, group)
else:
return self._medrecord.add_edges(edges)
edge_indices = self._medrecord.add_edges(edges)

if group is None:
return edge_indices

if not self.contains_group(group):
self.add_group(group)

self.add_edge_to_group(group, edge_indices)

return edge_indices

def add_edges_pandas(
self, edges: Union[PandasEdgeDataFrameInput, List[PandasEdgeDataFrameInput]]
self,
edges: Union[PandasEdgeDataFrameInput, List[PandasEdgeDataFrameInput]],
group: Optional[Group] = None,
) -> List[EdgeIndex]:
"""Adds edges to the MedRecord from one or more Pandas DataFrames.
"""Adds edges to the MedRecord from one or more Pandas DataFrames. Optionally assigns them to a group.
This method accepts either a single PandasEdgeDataFrameInput tuple or a list of
such tuples, each including a DataFrame and index columns for the source and
target nodes.
target nodes. If a group is specified, the edges are added to the group.
Args:
edges (Union[PandasEdgeDataFrameInput, List[PandasEdgeDataFrameInput]]):
A tuple or list of tuples, each including a DataFrame and index columns
for source and target nodes.
group (Optional[Group]): The name of the group to add the edges to. If not
specified, the edges are added to the MedRecord without a group.
Returns:
List[EdgeIndex]: A list of the edge indices added.
"""
return self.add_edges_polars(
[process_edges_dataframe(edges_df) for edges_df in edges]
if isinstance(edges, list)
else [process_edges_dataframe(edges)]
else [process_edges_dataframe(edges)],
group,
)

def add_edges_polars(
self,
edges: Union[PolarsEdgeDataFrameInput, List[PolarsEdgeDataFrameInput]],
group: Optional[Group] = None,
) -> List[EdgeIndex]:
"""Adds edges to the MedRecord from one or more Polars DataFrames.
"""Adds edges to the MedRecord from one or more Polars DataFrames. Optionally assigns them to a group.
This method accepts either a single PolarsEdgeDataFrameInput tuple or a list of
such tuples, each including a DataFrame and index columns for the source and
target nodes.
target nodes. If a group is specified, the edges are added to the group.
Args:
edges (Union[PolarsEdgeDataFrameInput, List[PolarsEdgeDataFrameInput]]):
A tuple or list of tuples, each including a DataFrame and index columns
for source and target nodes.
group (Optional[Group]): The name of the group to add the edges to. If not
specified, the edges are added to the MedRecord without a group.
Returns:
List[EdgeIndex]: A list of the edge indices added.
"""
return self._medrecord.add_edges_dataframes(
edge_indices = self._medrecord.add_edges_dataframes(
edges if isinstance(edges, list) else [edges]
)

if group is None:
return edge_indices

if not self.contains_group(group):
self.add_group(group)

self.add_edge_to_group(group, edge_indices)

return edge_indices

def add_group(
self,
group: Group,
Expand Down
Loading

0 comments on commit 3630bff

Please sign in to comment.