Merge pull request #82 from mariuzka/dev

Dev

pyproject.toml

@@ -108,7 +108,7 @@ convention = "google"
 [tool.ruff.lint.isort]
 force-single-line = true
 force-sort-within-sections = true
-lines-after-imports = -1
+known-first-party = ["src", "pop2net"]
 
 [build-system]
 requires = ["poetry-core"]

src/pop2net/inspector.py

@@ -40,10 +40,13 @@ def plot_bipartite_network(
         """Plots the two-mode network of agents and locations.
 
         Args:
-            node_color (str, optional): The node attribute that determines the
-                color of the nodes. If None, the node color represents whether
-                it is a location or an agent instance.
-            node_attrs (list | None, optional): A list of agent and location attributes that
+            agent_color (str, optional): The agent attribute that determines the
+                color of the agent nodes.
+            agent_attrs (list | None, optional): A list of agent attributes that
+                should be shown as node attributes in the network graph. Defaults to None.
+            location_color (str, optional): The location attribute that determines the
+                color of the location nodes.
+            location_attrs (list | None, optional): A list of location attributes that
                 should be shown as node attributes in the network graph. Defaults to None.
             edge_alpha (str, optional): The node attribute that determines the edges' transparency.
                 Defaults to "weight".
@@ -90,19 +93,19 @@ def plot_bipartite_network(
 
     def plot_agent_network(
         self,
-        node_color: str | None = "firebrick",
-        node_attrs: list | None = None,
+        agent_color: str | None = None,
+        agent_attrs: list | None = None,
         edge_alpha: str = "weight",
         edge_color: str = "black",
         include_0_weights: bool = True,
     ) -> None:
         """Plots the agent network.
 
         Args:
-            node_color (str, optional): The node attribute that determines the
+            agent_color (str, optional): The node attribute that determines the
                 color of the nodes. If None, the node color represents whether
                 it is a location or an agent instance.
-            node_attrs (list | None, optional): A list of agent and location attributes that
+            agent_attrs (list | None, optional): A list of agent attributes that
                 should be shown as node attributes in the network graph. Defaults to None.
             edge_alpha (str, optional): The node attribute that determines the edges' transparency.
                 Defaults to "weight".
@@ -111,26 +114,74 @@ def plot_agent_network(
             include_0_weights (bool, optional): Should edges with a weight of zero be included in
                 the plot? Defaults to True.
         """
-        if node_attrs is None:
-            node_attrs = ["type"]
+        if agent_attrs is None:
+            agent_attrs = ["type"]
         else:
-            node_attrs = list(node_attrs)
-            if "type" not in node_attrs:
-                node_attrs.append("type")
+            agent_attrs = list(agent_attrs)
+            if "type" not in agent_attrs:
+                agent_attrs.append("type")
+
+        if agent_color is not None and agent_color not in agent_attrs:
+            agent_attrs.append(agent_color)
 
         graph = self.model.export_agent_network(
-            node_attrs=node_attrs,
+            node_attrs=agent_attrs,
             include_0_weights=include_0_weights,
         )
 
         graph_layout = nx.drawing.spring_layout(graph)
         plot = BokehGraph(graph, width=400, height=400, hover_edges=True)
         plot.layout(layout=graph_layout)
         plot.draw(
-            node_color=node_color,
+            node_color="firebrick" if agent_color is None else agent_color,
+            edge_alpha=edge_alpha,
+            edge_color=edge_color,
+        )
+
+    def plot_networks(
+        self,
+        agent_attrs: list | None = None,
+        location_attrs: list | None = None,
+        agent_color: str | None = None,
+        location_color: str | None = None,
+        edge_alpha: str = "weight",
+        edge_color: str = "black",
+        include_0_weights: bool = True,
+    ) -> None:
+        """Plots the two-mode network and the agent network.
+
+        Args:
+            agent_color (str, optional): The agent attribute that determines the
+                color of the agent nodes.
+            agent_attrs (list | None, optional): A list of agent attributes that
+                should be shown as node attributes in the network graph. Defaults to None.
+            location_color (str, optional): The location attribute that determines the
+                color of the location nodes.
+            location_attrs (list | None, optional): A list of location attributes that
+                should be shown as node attributes in the network graph. Defaults to None.
+            edge_alpha (str, optional): The node attribute that determines the edges' transparency.
+                Defaults to "weight".
+            edge_color (str, optional): The node attribute that determines the edges' color.
+                Defaults to "black".
+            include_0_weights (bool): Should edges with a weight of zero be included in
+                the plot? Defaults to True.
+        """
+        self.plot_bipartite_network(
+            agent_color=agent_color,
+            agent_attrs=agent_attrs,
+            location_color=location_color,
+            location_attrs=location_attrs,
+            edge_color=edge_color,
             edge_alpha=edge_alpha,
+            # include_0_weights=include_0_weights,
+        )
+
+        self.plot_agent_network(
+            agent_color=agent_color,
+            agent_attrs=agent_attrs,
             edge_color=edge_color,
-            node_palette="random",
+            edge_alpha=edge_alpha,
+            include_0_weights=include_0_weights,
         )
 
     def eval_affiliations(self, return_data=False) -> tuple[pd.DataFrame, pd.DataFrame] | None:

tests/test_location/test_combinations/test_exact_n__agents_refine.py

@@ -28,7 +28,7 @@ def refine(self):
 
     inspector = p2n.NetworkInspector(model)
     inspector.plot_bipartite_network()
-    inspector.plot_agent_network(node_attrs=df.columns, node_color="status")
+    inspector.plot_agent_network(agent_attrs=df.columns, agent_color="status")
 
     assert len(model.agents) == 7
     assert len(model.locations) == 2
@@ -40,7 +40,3 @@ def refine(self):
     assert sum(agent.status == "A" for agent in model.locations[1].agents) == 1
     assert sum(agent.status == "B" for agent in model.locations[1].agents) == 1
     assert sum(agent.status == "C" for agent in model.locations[1].agents) == 1
-
-
-test_1()
-# %%

tests/test_location/test_combinations/test_exact_n_agents_filter.py

@@ -4,7 +4,6 @@
 import pop2net as p2n
 
 
-# %%
 def test_1():
     df = pd.DataFrame(
         {
@@ -24,7 +23,7 @@ def filter(self, agent):
     creator.create(df=df, location_classes=[TestLocation])
     inspector = p2n.NetworkInspector(model=model)
     inspector.plot_bipartite_network()
-    inspector.plot_agent_network(node_attrs=["status"])
+    inspector.plot_agent_network(agent_attrs=["status"])
 
     assert len(model.agents) == 4
     assert len(model.locations) == 2
@@ -47,14 +46,10 @@ def filter(self, agent):
     creator.create(df=df, location_classes=[TestLocation])
     inspector = p2n.NetworkInspector(model=model)
     inspector.plot_bipartite_network()
-    inspector.plot_agent_network(node_attrs=["status"])
+    inspector.plot_agent_network(agent_attrs=["status"])
 
     assert len(model.agents) == 4
     assert len(model.locations) == 1
     assert len(model.locations[0].agents) == 2
     assert all(agent.status == "A" for agent in model.locations[0].agents)
     assert all(not agent.locations for agent in model.agents if agent.status == "B")
-
-
-test_1()
-# %%

tests/test_location/test_combinations/test_exact_n_agents_nest.py

@@ -1,10 +1,7 @@
-# %%
 import pandas as pd
 
 import pop2net as p2n
 
-# %%
-
 
 def test_1():
     df = pd.DataFrame({"status": ["pupil", "pupil", "pupil", "pupil", "pupil"]})
@@ -46,18 +43,14 @@ class School(p2n.MagicLocation):
     creator.create(df=df, location_classes=[Classroom, School])
     inspector = p2n.NetworkInspector(model=model)
     inspector.plot_bipartite_network()
-    inspector.plot_agent_network(node_attrs=["status"])
+    inspector.plot_agent_network(agent_attrs=["status"])
 
     assert len(model.agents) == 5
     assert len(model.locations) == 2
     assert len(model.locations[0].agents) == 3
     assert len(model.locations[1].agents) == 5
 
 
-test_1()
-
-
-# %%
 def test_2():
     df = pd.DataFrame({"status": ["pupil", "pupil", "pupil", "pupil", "pupil", "pupil"]})
     model = p2n.Model()
@@ -77,7 +70,7 @@ class School(p2n.MagicLocation):
     creator.create(df=df, location_classes=[Classroom, School])
     inspector = p2n.NetworkInspector(model=model)
     inspector.plot_bipartite_network()
-    inspector.plot_agent_network(node_attrs=["status", "id"])
+    inspector.plot_agent_network(agent_attrs=["status", "id"])
 
     assert len(model.agents) == 6
     assert len(model.locations) == 4
@@ -86,6 +79,3 @@ class School(p2n.MagicLocation):
     assert len(model.locations[2].agents) == 2
     assert len(model.locations[3].agents) == 4
     assert model.locations[2].agents not in model.locations[3].agents
-
-
-test_2()

tests/test_location/test_combinations/test_exact_n_agents_split.py

@@ -1,10 +1,8 @@
-# %%
 import pandas as pd
 
 import pop2net as p2n
 
 
-# %%
 def test_1():
     df = pd.DataFrame(
         {
@@ -24,7 +22,7 @@ def split(self, agent):
     creator.create(df=df, location_classes=[TestLocation])
     inspector = p2n.NetworkInspector(model=model)
     inspector.plot_bipartite_network()
-    inspector.plot_agent_network(node_attrs=["status"])
+    inspector.plot_agent_network(agent_attrs=["status"])
 
     assert len(model.agents) == 6
     assert len(model.locations) == 4
@@ -50,7 +48,7 @@ def split(self, agent):
     creator.create(df=df, location_classes=[TestLocation])
     inspector = p2n.NetworkInspector(model=model)
     inspector.plot_bipartite_network()
-    inspector.plot_agent_network(node_attrs=["status"])
+    inspector.plot_agent_network(agent_attrs=["status"])
 
     assert len(model.agents) == 6
     assert len(model.locations) == 2
@@ -60,7 +58,3 @@ def split(self, agent):
     assert all(agent.status == "B" for agent in model.locations[1].agents)
     assert sum(not agent.locations for agent in model.agents if agent.status == "B") == 1
     assert sum(not agent.locations for agent in model.agents if agent.status == "A") == 1
-
-
-test_1()
-# %%

tests/test_location/test_combinations/test_exact_n_agents_weight.py

@@ -1,10 +1,7 @@
-# %%
 import pandas as pd
 
 import pop2net as p2n
 
-# %%
-
 
 def test_1():
     df = pd.DataFrame(
@@ -31,6 +28,3 @@ def weight(self, agent):
     assert len(model.locations[0].agents) == 2
     assert all(agent.status == "A" for agent in model.locations[0].agents)
     assert all(model.locations[0].get_weight(agent) == 1 for agent in model.locations[0].agents)
-
-
-test_1()

tests/test_location/test_combinations/test_exct_n_agents_melt.py

@@ -1,14 +1,12 @@
-# %%
 import pandas as pd
 
 import pop2net as p2n
 
-# %%
-
-
 # Testen: Wie verhält es sich wenn n_agents vorgegeben ist aber noch lehrer übrig sind
 # wie verhält es sich, wenn in der Ziel location n_agent und exact true ist, aber in den unteren nichts gesetz ist
 # was passiert wenn ich false und true in den Meltlocation mixe?
+
+
 def test_1():
     df = pd.DataFrame(
         {
@@ -41,7 +39,7 @@ def melt(self):
     creator.create_locations(location_classes=[LocAB])
     inspector = p2n.NetworkInspector(model)
     inspector.plot_bipartite_network()
-    inspector.plot_agent_network(node_attrs=["status"])
+    inspector.plot_agent_network(agent_attrs=["status"])
 
     assert len(model.agents) == 6
     assert len(model.locations) == 2
@@ -79,14 +77,10 @@ def melt(self):
     creator.create_locations(location_classes=[LocAB])
     inspector = p2n.NetworkInspector(model)
     inspector.plot_bipartite_network()
-    inspector.plot_agent_network(node_attrs=["status"])
+    inspector.plot_agent_network(agent_attrs=["status"])
 
     assert len(model.agents) == 6
     assert len(model.locations) == 1
     assert sum(True for agent in model.locations[0].agents if agent.status == "A") == 2
     assert sum(True for agent in model.locations[0].agents if agent.status == "B") == 2
     assert sum(True for agent in model.agents if not agent.locations) == 2
-
-
-test_1()
-# %%

tests/test_location/test_combinations/test_filter_melt.py

@@ -1,10 +1,8 @@
-# %%
 import pandas as pd
 
 import pop2net as p2n
 
 
-# %%
 def test_1():
     model = p2n.Model()
     creator = p2n.Creator(model=model)
@@ -35,10 +33,6 @@ def melt(self):
     assert sum(agent.status == "pupil" for agent in model.locations[0].agents) == 4
 
 
-test_1()
-
-
-# %%
 def test_2():
     model = p2n.Model()
     creator = p2n.Creator(model=model)
@@ -73,7 +67,3 @@ def melt(self):
     assert sum(agent.status == "pupil" for agent in model.locations[0].agents) == 2
     assert sum(agent.status == "teacher" for agent in model.locations[0].agents) == 1
     assert sum(not agent.locations for agent in model.agents) == 3
-
-
-test_2()
-# %%