fix(connec): fix issue with line_color display under color scale

When `scale_method` is set to `color` or `width_color` and `line_color` is specified,
the color was not applied correctly. This fix ensures the color is displayed as intended.

---

修复(connec): 修复 color scale 下 line_color 显示异常问题

当 `scale_method` 设置为 `color` 或 `width_color` 且同时指定了 `line_color` 时，
颜色未能正确应用。该问题已修复，确保颜色显示符合预期。

Author: RicardoRyn

src/plotfig/brain_connection.py

@@ -1,16 +1,18 @@
-import os
-import os.path as op
 import datetime
+from pathlib import Path
+from typing import Literal
+from collections.abc import Sequence
+
 import numpy as np
+import numpy.typing as npt
 import nibabel as nib
-from scipy.ndimage import center_of_mass
 import plotly.graph_objects as go
 import plotly.io as pio
-import matplotlib.colors as mcolors
-import matplotlib.cm as cm
+from matplotlib.colors import LinearSegmentedColormap, to_hex
+from scipy.ndimage import center_of_mass
 from tqdm import tqdm
-from pathlib import Path
-import numpy.typing as npt
+
+from loguru import logger
 
 Num = int | float
 
@@ -19,15 +21,30 @@
     "save_brain_connection_frames",
 ]
 
-def _load_surface(file: str | Path):
-    '''加载 .surf.gii 文件，提取顶点和面'''
-    gii = nib.load(file)
-    vertices = gii.darrays[0].data
-    faces = gii.darrays[1].data
-    return vertices, faces
+
+def _validate_connectome(connectome):
+    """检测数据是否为有效的对称且对角线为0的连接矩阵"""
+    # 1. 判断是否二维方阵
+    if connectome.ndim != 2 or connectome.shape[0] != connectome.shape[1]:
+        raise ValueError("connectome 必须是二维方阵")
+    # 2. 判断是否对称矩阵
+    if not np.allclose(connectome, connectome.T, atol=1e-8):
+        raise ValueError("connectome 必须是对称矩阵")
+    # 3. 判断对角线是否全为0
+    if not np.allclose(np.diag(connectome), 0, atol=1e-8):
+        raise ValueError("connectome 对角线必须全部为0")
+    # 4. 判断是否全0矩阵，警告但不抛异常
+    if np.allclose(connectome, 0, atol=1e-8):
+        logger.warning("connectome 矩阵所有元素均为0，可能没有有效连接数据")
+
+
+def _load_surface(file):
+    """加载 .surf.gii 文件，提取顶点和面"""
+    return nib.load(file).darrays[0].data, nib.load(file).darrays[1].data
+
 
 def _create_mesh(vertices, faces, name):
-    '''	创建 plotly 的 Mesh3d 图层'''
+    """创建 plotly 的 Mesh3d 图层"""
     return go.Mesh3d(
         x=vertices[:, 0],
         y=vertices[:, 1],
@@ -42,29 +59,33 @@ def _create_mesh(vertices, faces, name):
         name=name,
     )
 
+
 def _get_node_indices(connectome, show_all_nodes):
-    '''	判断哪些节点需要显示'''
+    """判断是否显示无任何连接的节点"""
     if not show_all_nodes:
         row_is_zero = np.any(connectome != 0, axis=1)
         return np.where(row_is_zero)[0]
     else:
         return np.arange(connectome.shape[0])
 
-def _get_centroids_real(niigz_file: str | Path):
-    '''读取 NIfTI 图集并计算ROI质心'''
-    img = nib.load(niigz_file)
-    atlas_data = img.get_fdata()
-    affine = img.affine
 
+def _get_centroids_real(niigz_file):
+    """读取 NIfTI 图集并计算ROI质心"""
+    atlas_data = nib.load(niigz_file).get_fdata()
+    affine = nib.load(niigz_file).affine
     roi_labels = np.unique(atlas_data)
     roi_labels = roi_labels[roi_labels != 0]
-
-    centroids_voxel = [center_of_mass((atlas_data == label).astype(int)) for label in roi_labels]
+    centroids_voxel = [
+        center_of_mass((atlas_data == label).astype(int)) for label in roi_labels
+    ]
     centroids_real = [np.dot(affine, [*coord, 1])[:3] for coord in centroids_voxel]
     return np.array(centroids_real)
 
-def _add_nodes_to_fig(fig, centroids_real, node_indices, nodes_name, nodes_size, nodes_color):
-    '''将节点（球）添加到图中'''
+
+def _add_nodes_to_fig(
+    fig, centroids_real, node_indices, nodes_name, nodes_size, nodes_color
+):
+    """将节点（球）添加到图中"""
     for i in node_indices:
         fig.add_trace(
             go.Scatter3d(
@@ -74,7 +95,7 @@ def _add_nodes_to_fig(fig, centroids_real, node_indices, nodes_name, nodes_size,
                 mode="markers+text",
                 marker={
                     "size": nodes_size[i],
-                    "color": nodes_color,
+                    "color": nodes_color[i],
                     "colorscale": "Rainbow",
                     "opacity": 0.8,
                     "line": {"width": 2, "color": "black"},
@@ -85,8 +106,25 @@ def _add_nodes_to_fig(fig, centroids_real, node_indices, nodes_name, nodes_size,
             )
         )
 
-def _add_edges_to_fig(fig, connectome, centroids_real, nodes_name, scale_method, line_width, line_color="#ff0000"):
-    '''将连接线绘制到图中'''
+
+def _add_edges_to_fig(
+    fig,
+    connectome,
+    centroids_real,
+    nodes_name,
+    scale_method,
+    line_width,
+    line_color,
+):
+    """将连接线绘制到图中"""
+
+    def _get_gradient_color(value, color):
+        """获取渐变色"""
+        assert 0 <= value <= 1, "value 必须在0和1之间"
+        cmap = LinearSegmentedColormap.from_list("grad_cmap", ["white", color])
+        rgba = cmap(value)
+        return to_hex(rgba[:3])
+
     nodes_num = connectome.shape[0]
     if np.all(connectome == 0):
         return
@@ -100,41 +138,50 @@ def _add_edges_to_fig(fig, connectome, centroids_real, nodes_name, scale_method,
                 continue
 
             match scale_method:
+                case "":
+                    each_line_color = line_color if value > 0 else "#0000ff"
+                    each_line_width = line_width
                 case "width":
                     each_line_color = line_color if value > 0 else "#0000ff"
                     each_line_width = abs(value / max_strength) * line_width
                 case "color":
                     norm_value = value / max_strength
-                    each_line_color = mcolors.to_hex(cm.bwr(mcolors.Normalize(vmin=-1, vmax=1)(norm_value)))
+                    each_line_color = _get_gradient_color(norm_value, line_color)
                     each_line_width = line_width
                 case "width_color" | "color_width":
                     norm_value = value / max_strength
                     each_line_width = abs(norm_value) * line_width
-                    each_line_color = mcolors.to_hex(cm.bwr(mcolors.Normalize(vmin=-1, vmax=1)(norm_value)))
-                case "":
-                    each_line_color = "#ff0000" if value > 0 else "#0000ff"
-                    each_line_width = line_width
+                    each_line_color = _get_gradient_color(norm_value, line_color)
                 case _:
-                    raise ValueError("scale_method must be '', 'width', 'color', 'width_color', or 'color_width'")
+                    raise ValueError(
+                        "scale_method 必须为 '', 'width', 'color', 'width_color', or 'color_width'中的一种"
+                    )
 
-            connection_line = np.array([centroids_real[i], centroids_real[j], [None] * 3])
+            connection_line = np.array(
+                [centroids_real[i], centroids_real[j], [None] * 3]
+            )
             fig.add_trace(
                 go.Scatter3d(
                     x=connection_line[:, 0],
                     y=connection_line[:, 1],
                     z=connection_line[:, 2],
                     mode="lines",
                     line={"color": each_line_color, "width": each_line_width},
-                    hoverinfo="none",
+                    hoverinfo="name",
                     name=f"{nodes_name[i]}-{nodes_name[j]}",
                 )
             )
 
+
 def _finalize_figure(fig):
-    '''调整图形布局与视觉样式'''
+    """调整图形布局与视觉样式"""
     fig.update_traces(
         selector={"mode": "markers"},
-        marker={"size": 10, "colorscale": "Viridis", "line": {"width": 3, "color": "black"}},
+        marker={
+            "size": 10,
+            "colorscale": "Viridis",
+            "line": {"width": 3, "color": "black"},
+        },
     )
     fig.update_layout(
         title="Connection",
@@ -147,80 +194,116 @@ def _finalize_figure(fig):
         margin={"l": 0, "r": 0, "b": 0, "t": 30},
     )
 
+
 def plot_brain_connection_figure(
     connectome: npt.NDArray,
     lh_surfgii_file: str | Path,
     rh_surfgii_file: str | Path,
     niigz_file: str | Path,
+    output_file: str | Path | None = None,
+    show_all_nodes: bool = False,
+    nodes_size: Sequence[Num] | npt.NDArray | None = None,
     nodes_name: list[str] | None = None,
-    nodes_size=None,
     nodes_color: list[str] | None = None,
-    output_file: str | Path | None = None,
-    scale_method: str = "",
+    scale_method: Literal["", "width", "color", "width_color", "color_width"] = "",
     line_width: Num = 10,
-    show_all_nodes: bool = False,
-    line_color: str = "#ff0000",
-) -> None:
+    line_color: str = "red",
+) -> go.Figure:
     """绘制大脑连接图，保存在指定的html文件中
 
     Args:
-        connectome (npt.NDArray): 连接矩阵
-        lh_surfgii_file (str | Path): 左脑surf.gii文件.
-        rh_surfgii_file (str | Path): 右脑surf.gii文件.
-        niigz_file (str | Path): 图集nii文件.
-        nodes_name (List[str] | None, optional): 节点名称. Defaults to None.
-        nodes_size (Num, optional): 节点大小. Defaults to 5.
-        nodes_color (List[str] | None, optional): 节点颜色. Defaults to None.
-        output_file (str | Path | None, optional): 保存的完整路径及文件名. Defaults to None.
-        scale_method (str, optional): 连接scale的形式. Defaults to "".
-        line_width (Num, optional): 连接粗细. Defaults to 10.
-
-    Raises:
-        ValueError: 参数参数取值不合法时抛出.
+        connectome (npt.NDArray):
+            大脑连接矩阵，形状为 (n, n)，其中 n 是脑区数量。
+            矩阵中的值表示脑区之间的连接强度，正值表示正相关连接，负值表示负相关连接，0表示无连接。
+        lh_surfgii_file (str | Path):
+            左半脑表面几何文件路径 (.surf.gii 格式)，用于绘制左半脑表面
+        rh_surfgii_file (str | Path):
+            右半脑表面几何文件路径 (.surf.gii 格式)，用于绘制右半脑表面
+        niigz_file (str | Path):
+            NIfTI格式的脑区图谱文件路径 (.nii.gz 格式)，用于定位脑区节点的三维坐标
+        output_file (str | Path | None, optional):
+            输出HTML文件路径。如果未指定，则使用当前时间戳生成文件名。默认为None
+        show_all_nodes (bool, optional):
+            是否显示所有脑区节点。如果为False，则只显示有连接的节点。默认为False
+        nodes_size (Sequence[Num] | npt.NDArray | None, optional):
+            每个节点的大小，长度应与脑区数量一致。默认为None，即所有节点大小为5
+        nodes_name (list[str] | None, optional):
+            每个节点的名称标签，长度应与脑区数量一致。默认为None，即不显示名称
+        nodes_color (list[str] | None, optional):
+            每个节点的颜色，长度应与脑区数量一致。默认为None，即所有节点为白色
+        scale_method (Literal["", "width", "color", "width_color", "color_width"], optional):
+            连接线的缩放方法:
+            - "" : 所有连接线宽度和颜色固定
+            - "width" : 根据连接强度调整线宽，正连接为红色，负连接为蓝色
+            - "color" : 根据连接强度调整颜色(使用蓝白红颜色映射)，线宽固定
+            - "width_color" or "color_width" : 同时根据连接强度调整线宽和颜色
+            默认为 ""
+        line_width (Num, optional):
+            连接线的基本宽度。当scale_method包含"width"时，此值作为最大宽度参考。默认为10
+        line_color (str, optional):
+            连接线的基本颜色。当scale_method不包含"color"时生效。默认为"#ff0000"(红色)
+
+    Returns:
+        go.Figure: Plotly图形对象，包含绘制的大脑连接图
     """
+    _validate_connectome(connectome)
+
+    if np.any(connectome < 0):
+        logger.warning(
+            "由于 connectome 存在负值，连线颜色无法自定义，只能正值显示红色，负值显示蓝色"
+        )
+        line_color = "#ff0000"
+
     nodes_num = connectome.shape[0]
-    if nodes_name is None:
-        nodes_name = [f"ROI-{i}" for i in range(nodes_num)]
-    if nodes_color is None:
-        nodes_color = ["white"] * nodes_num
-    if nodes_size is None:
-        nodes_size = [5] * nodes_num
+    nodes_name = nodes_name or [""] * nodes_num
+    nodes_color = nodes_color or ["white"] * nodes_num
+    nodes_size = nodes_size or [5] * nodes_num
+
     if output_file is None:
         timestamp = datetime.datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
-        output_file = op.join(f"./{timestamp}.html")
-        print(f"未指定保存路径，默认保存为：{output_file}")
+        output_file = Path(f"{timestamp}.html")
+        logger.info(f"未指定保存路径，默认保存在当前文件夹下的{output_file}中。")
 
     node_indices = _get_node_indices(connectome, show_all_nodes)
     vertices_L, faces_L = _load_surface(lh_surfgii_file)
     vertices_R, faces_R = _load_surface(rh_surfgii_file)
 
     mesh_L = _create_mesh(vertices_L, faces_L, "Left Hemisphere")
     mesh_R = _create_mesh(vertices_R, faces_R, "Right Hemisphere")
+
     fig = go.Figure(data=[mesh_L, mesh_R])
 
     centroids_real = _get_centroids_real(niigz_file)
-    _add_nodes_to_fig(fig, centroids_real, node_indices, nodes_name, nodes_size, nodes_color)
-    _add_edges_to_fig(fig, connectome, centroids_real, nodes_name, scale_method, line_width, line_color)
+    _add_nodes_to_fig(
+        fig, centroids_real, node_indices, nodes_name, nodes_size, nodes_color
+    )
+    _add_edges_to_fig(
+        fig,
+        connectome,
+        centroids_real,
+        nodes_name,
+        scale_method,
+        line_width,
+        line_color,
+    )
     _finalize_figure(fig)
 
     fig.write_html(output_file)
     return fig
 
 
 def save_brain_connection_frames(
-    fig: go.Figure,
-    output_dir: str,
-    n_frames: int = 36
+    fig: go.Figure, output_dir: str | Path, n_frames: int = 36
 ) -> None:
     """
-    生成不同角度的静态图片帧，用于制作旋转大脑连接图的 GIF 或视频。
+    生成不同角度的静态图片帧，可用于制作旋转大脑连接图的 GIF。
 
     Args:
         fig (go.Figure): Plotly 的 Figure 对象，包含大脑表面和连接图。
-        output_dir (str): 图片保存的文件夹路径，会自动创建文件夹。
+        output_dir (str): 图片保存的文件夹路径，若文件夹不存在则自动创建。
         n_frames (int, optional): 旋转帧的数量。默认 36，即每 10 度一帧。
     """
-    os.makedirs(output_dir, exist_ok=True)
+    Path(output_dir).mkdir(parents=True, exist_ok=True)
     angles = np.linspace(0, 360, n_frames, endpoint=False)
     for i, angle in tqdm(enumerate(angles), total=len(angles)):
         camera = dict(
@@ -230,12 +313,4 @@ def save_brain_connection_frames(
         )
         fig.update_layout(scene_camera=camera)
         pio.write_image(fig, f"{output_dir}/frame_{i:03d}.png", width=800, height=800)
-    print(f"保存了 {n_frames} 张图片在 {output_dir}")
-
-
-def main():
-    pass
-
-
-if __name__ == "__main__":
-    main()
+    logger.info(f"保存了 {n_frames} 张图片在 {output_dir}")