from __future__ import annotations import pytest import numpy as np import pandas as pd from datashader.bundling import directly_connect_edges, hammer_bundle from datashader.layout import circular_layout, forceatlas2_layout, random_layout skimage = pytest.importorskip("skimage") @pytest.fixture def nodes(): # Four nodes arranged at the corners of a 200x200 square with one node # at the center nodes_df = pd.DataFrame({'id': np.arange(5), 'x': [0, -100, 100, -100, 100], 'y': [0, 100, 100, -100, -100]}) nodes_df.set_index('id') return nodes_df @pytest.fixture def edges(): # Four edges originating from the center node and connected to each # corner edges_df = pd.DataFrame({'id': np.arange(4), 'source': np.zeros(4, dtype=int), 'target': np.arange(1, 5)}) edges_df.set_index('id') return edges_df @pytest.fixture def weighted_edges(): # Four weighted edges originating from the center node and connected # to each corner edges_df = pd.DataFrame({'id': np.arange(4), 'source': np.zeros(4, dtype=int), 'target': np.arange(1, 5), 'weight': np.ones(4)}) edges_df.set_index('id') return edges_df def test_immutable_nodes(nodes, edges): # Expect nodes to remain immutable after any bundling operation original = nodes.copy() directly_connect_edges(nodes, edges) assert original.equals(nodes) @pytest.mark.parametrize('bundle', [directly_connect_edges, hammer_bundle]) def test_renamed_columns(nodes, weighted_edges, bundle): nodes = nodes.rename(columns={'x': 'xx', 'y': 'yy'}) edges = weighted_edges.rename(columns={'source': 'src', 'target': 'dst', 'weight': 'w'}) df = bundle(nodes, edges, x='xx', y='yy', source='src', target='dst', weight='w') assert 'xx' in df and 'x' not in df assert 'yy' in df and 'y' not in df assert 'w' in df and 'weight' not in df @pytest.mark.parametrize('bundle', [directly_connect_edges, hammer_bundle]) @pytest.mark.parametrize('layout', [random_layout, circular_layout, forceatlas2_layout]) def test_same_path_endpoints(layout, bundle): # Expect path endpoints to match original edge source/target edges = pd.DataFrame({'id': [0], 'source': [0], 'target': [1]}).set_index('id') nodes = pd.DataFrame({'id': np.unique(edges.values)}).set_index('id') node_positions = layout(nodes, edges) bundled = bundle(node_positions, edges) source, target = edges.iloc[0] expected_source = node_positions.loc[source] expected_target = node_positions.loc[target] actual_source = bundled.loc[0] actual_target = bundled.loc[len(bundled)-2] assert np.allclose(expected_source, actual_source) assert np.allclose(expected_target, actual_target) @pytest.mark.parametrize("include_edge_id", [True, False]) def test_directly_connect_with_weights(nodes, weighted_edges, include_edge_id): # Expect four lines starting at center (0.5, 0.5) and terminating # at a different corner and NaN data = pd.DataFrame({'edge_id': [1.0, 1.0, np.nan, 2.0, 2.0, np.nan, 3.0, 3.0, np.nan, 4.0, 4.0, np.nan], 'x': [0.0, -100.0, np.nan, 0.0, 100.0, np.nan, 0.0, -100.0, np.nan, 0.0, 100.0, np.nan], 'y': [0.0, 100.0, np.nan, 0.0, 100.0, np.nan, 0.0, -100.0, np.nan, 0.0, -100.0, np.nan]}) columns = ['edge_id', 'x', 'y'] if include_edge_id else ['x', 'y'] expected = pd.DataFrame(data, columns=columns) given = directly_connect_edges(nodes, weighted_edges, include_edge_id=include_edge_id) assert given.equals(expected) @pytest.mark.parametrize("include_edge_id", [True, False]) def test_directly_connect_without_weights(nodes, edges, include_edge_id): # Expect four lines starting at center (0.5, 0.5) and terminating # at a different corner and NaN data = pd.DataFrame({'edge_id': [1.0, 1.0, np.nan, 2.0, 2.0, np.nan, 3.0, 3.0, np.nan, 4.0, 4.0, np.nan], 'x': [0.0, -100.0, np.nan, 0.0, 100.0, np.nan, 0.0, -100.0, np.nan, 0.0, 100.0, np.nan], 'y': [0.0, 100.0, np.nan, 0.0, 100.0, np.nan, 0.0, -100.0, np.nan, 0.0, -100.0, np.nan]}) columns = ['edge_id', 'x', 'y'] if include_edge_id else ['x', 'y'] expected = pd.DataFrame(data, columns=columns) given = directly_connect_edges(nodes, edges, include_edge_id=include_edge_id) assert given.equals(expected) @pytest.mark.parametrize("include_edge_id", [True, False]) def test_hammer_bundle_with_weights(nodes, weighted_edges, include_edge_id): # Expect four lines starting at center (0.0, 0.0) and terminating # with NaN data = pd.DataFrame({'edge_id': [1.0, np.nan, 2.0, np.nan, 3.0, np.nan, 4.0, np.nan], 'x': [0.0, np.nan, 0.0, np.nan, 0.0, np.nan, 0.0, np.nan], 'y': [0.0, np.nan, 0.0, np.nan, 0.0, np.nan, 0.0, np.nan], 'weight': [1.0, np.nan, 1.0, np.nan, 1.0, np.nan, 1.0, np.nan]}) columns = ['edge_id', 'x', 'y', 'weight'] if include_edge_id else ['x', 'y', 'weight'] expected = pd.DataFrame(data, columns=columns) df = hammer_bundle(nodes, weighted_edges, include_edge_id=include_edge_id) starts = df[(df.x == 0.0) & (df.y == 0.0)] ends = df[df.isnull().any(axis=1)] given = pd.concat([starts, ends]) given.sort_index(inplace=True) given.reset_index(drop=True, inplace=True) assert given.equals(expected) @pytest.mark.parametrize("include_edge_id", [True, False]) def test_hammer_bundle_without_weights(nodes, edges, include_edge_id): # Expect four lines starting at center (0.0, 0.0) and terminating # with NaN data = pd.DataFrame({'edge_id': [1.0, np.nan, 2.0, np.nan, 3.0, np.nan, 4.0, np.nan], 'x': [0.0, np.nan, 0.0, np.nan, 0.0, np.nan, 0.0, np.nan], 'y': [0.0, np.nan, 0.0, np.nan, 0.0, np.nan, 0.0, np.nan]}) columns = ['edge_id', 'x', 'y'] if include_edge_id else ['x', 'y'] expected = pd.DataFrame(data, columns=columns) df = hammer_bundle(nodes, edges, include_edge_id=include_edge_id) starts = df[(df.x == 0.0) & (df.y == 0.0)] ends = df[df.isnull().any(axis=1)] given = pd.concat([starts, ends]) given.sort_index(inplace=True) given.reset_index(drop=True, inplace=True) assert given.equals(expected)