diff --git a/prospr/_version.py b/prospr/_version.py index 530b385..d7e1590 100644 --- a/prospr/_version.py +++ b/prospr/_version.py @@ -1 +1 @@ -__version__ = "1.2.14" +__version__ = "1.2.15" diff --git a/prospr/visualize.py b/prospr/visualize.py index ea67659..514105a 100644 --- a/prospr/visualize.py +++ b/prospr/visualize.py @@ -26,7 +26,7 @@ # Global colors. COLOR_H = "tab:blue" COLOR_P = "orange" -COLOR_ANNOTATE = "#00ce00" +COLOR_ANNOTATE = "#ff00ff" COLOR_CONTACT = "indianred" COLOR_CHAIN = "black" @@ -43,13 +43,16 @@ def _plot_aminos_2d_basic( :param float markersize: Size of the markers. :param bool annotate_first: Annotate first amino acid. """ + # Plot protein chain. ax.plot( df["x"], df["y"], color=COLOR_CHAIN, alpha=0.65, zorder=1, lw=linewidth ) + + # Plot all aminos except first as it might be annotated. sns.scatterplot( x="x", y="y", - data=df, + data=df.iloc[1:], hue="Type", hue_order=["H", "P"], style="Type", @@ -60,26 +63,26 @@ def _plot_aminos_2d_basic( ax=ax, ) - # Plot first point with a green edge color. + # Plot the first amino acid as a diamond when annotating it. + if df.iloc[0]["Type"] == "H": + facecolor = COLOR_H + marker_first = "o" + elif df.iloc[0]["Type"] == "P": + facecolor = COLOR_P + marker_first = "s" + if annotate_first: - if df.iloc[0]["Type"] == "H": - ax.scatter( - df.iloc[0]["x"], - df.iloc[0]["y"], - marker="o", - fc=COLOR_ANNOTATE, - s=markersize, - zorder=2, - ) - else: - ax.scatter( - df.iloc[0]["x"], - df.iloc[0]["y"], - marker="s", - fc=COLOR_ANNOTATE, - s=markersize, - zorder=2, - ) + facecolor = COLOR_ANNOTATE + marker_first = "D" + + ax.scatter( + df.iloc[0]["x"], + df.iloc[0]["y"], + fc=facecolor, + marker=marker_first, + s=markersize, + zorder=2, + ) # Plot dotted lines between the aminos that increase the stability. pairs = get_scoring_pairs(protein) @@ -115,11 +118,7 @@ def _plot_aminos_2d_paper( if markersize < PAPER_MARKERSIZE: edgewidth *= 2 - # Split dataframe on amino acid type. - df_H = df.loc[df["Type"] == "H"] - df_P = df.loc[df["Type"] == "P"] - - # Plot protein. + # Plot protein chain. ax.plot( df["x"], df["y"], @@ -128,10 +127,12 @@ def _plot_aminos_2d_paper( linewidth=linewidth, zorder=1, ) + + # Plot everything but the first amino acid as it might be annotated. sns.scatterplot( x="x", y="y", - data=df_H, + data=df.iloc[1:].loc[df["Type"] == "H"], marker="o", facecolor=COLOR_H, edgecolor=COLOR_H, @@ -143,7 +144,7 @@ def _plot_aminos_2d_paper( sns.scatterplot( x="x", y="y", - data=df_P, + data=df.iloc[1:].loc[df["Type"] == "P"], marker="o", facecolor="white", edgecolor=COLOR_P, @@ -154,31 +155,29 @@ def _plot_aminos_2d_paper( label="P", ) - # Plot first point with a green edge color. + # Plot the first amino acid as a diamond when annotating it. + if df.iloc[0]["Type"] == "H": + edgecolor = COLOR_H + facecolor = COLOR_H + elif df.iloc[0]["Type"] == "P": + edgecolor = COLOR_P + facecolor = "white" + marker_first = "o" + if annotate_first: - # Plot first point with a different color. - if df.iloc[0]["Type"] == "H": - ax.scatter( - df.iloc[0]["x"], - df.iloc[0]["y"], - marker="o", - fc=COLOR_H, - ec=COLOR_ANNOTATE, - lw=edgewidth, - s=markersize, - zorder=2, - ) - else: - ax.scatter( - df.iloc[0]["x"], - df.iloc[0]["y"], - marker="o", - fc="white", - ec=COLOR_ANNOTATE, - lw=edgewidth, - s=markersize, - zorder=2, - ) + marker_first = "D" + edgecolor = COLOR_ANNOTATE + + ax.scatter( + df.iloc[0]["x"], + df.iloc[0]["y"], + marker=marker_first, + fc=facecolor, + ec=edgecolor, + lw=edgewidth, + s=markersize, + zorder=2, + ) # Plot dotted lines between the aminos that increase the stability. pairs = get_scoring_pairs(protein) @@ -210,10 +209,7 @@ def _plot_aminos_3d_basic( :param float markersize: Size of the markers. :param bool annotate_first: Annotate first amino acid. """ - # Split dataframe on amino acid type. - df_H = df.loc[df["Type"] == "H"] - df_P = df.loc[df["Type"] == "P"] - + # Plot protein chain. ax.plot( df["x"], df["y"], @@ -224,11 +220,11 @@ def _plot_aminos_3d_basic( lw=linewidth, ) - # Plot the aminos connected with an opaque line. + # Plot all except the first amino acid, as it might be annotated. ax.scatter( - df_H["x"], - df_H["y"], - df_H["z"], + df.iloc[1:].loc[df["Type"] == "H"]["x"], + df.iloc[1:].loc[df["Type"] == "H"]["y"], + df.iloc[1:].loc[df["Type"] == "H"]["z"], c=COLOR_H, marker="o", depthshade=False, @@ -236,9 +232,9 @@ def _plot_aminos_3d_basic( label="H", ) ax.scatter( - df_P["x"], - df_P["y"], - df_P["z"], + df.iloc[1:].loc[df["Type"] == "P"]["x"], + df.iloc[1:].loc[df["Type"] == "P"]["y"], + df.iloc[1:].loc[df["Type"] == "P"]["z"], c=COLOR_P, marker="s", depthshade=False, @@ -246,6 +242,27 @@ def _plot_aminos_3d_basic( label="P", ) + # Plot the first amino acid as a diamond when annotating it. + if df.iloc[0]["Type"] == "H": + facecolor = COLOR_H + elif df.iloc[0]["Type"] == "P": + facecolor = "white" + marker_first = "o" + + if annotate_first: + facecolor = COLOR_ANNOTATE + marker_first = "D" + + ax.scatter( + df.iloc[0]["x"], + df.iloc[0]["y"], + df.iloc[0]["z"], + fc=facecolor, + marker=marker_first, + s=markersize, + zorder=2, + ) + # Plot dotted lines between the aminos that increase the stability. pairs = get_scoring_pairs(protein) @@ -278,11 +295,7 @@ def _plot_aminos_3d_paper( if markersize < PAPER_MARKERSIZE: edgewidth *= 2 - # Split dataframe on amino acid type. - df_H = df.loc[df["Type"] == "H"] - df_P = df.loc[df["Type"] == "P"] - - # Plot protein. + # Plot protein chain. ax.plot( df["x"], df["y"], @@ -293,10 +306,11 @@ def _plot_aminos_3d_paper( zorder=1, ) + # Plot everything but the first amino acid as it might be annotated. ax.scatter( - df_H["x"], - df_H["y"], - df_H["z"], + df.iloc[1:].loc[df["Type"] == "H"]["x"], + df.iloc[1:].loc[df["Type"] == "H"]["y"], + df.iloc[1:].loc[df["Type"] == "H"]["z"], marker="o", edgecolor=COLOR_H, s=markersize, @@ -305,9 +319,9 @@ def _plot_aminos_3d_paper( depthshade=False, ) ax.scatter( - df_P["x"], - df_P["y"], - df_P["z"], + df.iloc[1:].loc[df["Type"] == "P"]["x"], + df.iloc[1:].loc[df["Type"] == "P"]["y"], + df.iloc[1:].loc[df["Type"] == "P"]["z"], marker="o", facecolor="white", edgecolor=COLOR_P, @@ -318,32 +332,30 @@ def _plot_aminos_3d_paper( depthshade=False, ) - # Plot first point with a green edge color. + # Plot the first amino acid as a diamond when annotating it. + if df.iloc[0]["Type"] == "H": + edgecolor = COLOR_H + facecolor = COLOR_H + if df.iloc[0]["Type"] == "P": + edgecolor = COLOR_P + facecolor = "white" + marker_first = "o" + if annotate_first: - if df.iloc[0]["Type"] == "H": - ax.scatter( - df.iloc[0]["x"], - df.iloc[0]["y"], - df.iloc[0]["z"], - marker="o", - fc=COLOR_H, - ec=COLOR_ANNOTATE, - lw=edgewidth, - s=markersize, - zorder=2, - ) - else: - ax.scatter( - df.iloc[0]["x"], - df.iloc[0]["y"], - df.iloc[0]["z"], - marker="o", - fc="white", - ec=COLOR_ANNOTATE, - lw=edgewidth, - s=markersize, - zorder=2, - ) + marker_first = "D" + edgecolor = COLOR_ANNOTATE + + ax.scatter( + df.iloc[0]["x"], + df.iloc[0]["y"], + df.iloc[0]["z"], + marker=marker_first, + fc=facecolor, + ec=edgecolor, + lw=edgewidth, + s=markersize, + zorder=2, + ) # Plot dotted lines between the aminos that increase the stability. pairs = get_scoring_pairs(protein) @@ -483,7 +495,7 @@ def plot_protein( linestyle=":", linewidth=linewidth, alpha=0.9, - label="Bond", + label="Contact", ) handles.append(score_patch) labels.append(score_patch.get_label()) diff --git a/tests/visualize/test_visualization.py b/tests/visualize/test_visualization.py index a547701..e844932 100644 --- a/tests/visualize/test_visualization.py +++ b/tests/visualize/test_visualization.py @@ -8,10 +8,6 @@ def test_2d_visualization(self, show=False): """ Test the 2D plotting in all different styles. """ - p = Protein("HPHPPPPPPP", dim=2, model="HP") - p.set_hash([-1, -1, 2, 2, 2, 2, 2, 2, 2]) - plot_protein(p, style="paper", show=False) - p = Protein("HPPHPPHPPH", dim=2, model="HP") p.set_hash([-1, -2, 1, -2, 1, 2, 1, 2, -1]) plot_protein(p, style="paper", show=False) @@ -50,13 +46,25 @@ def test_2d_visualization(self, show=False): p.set_hash([-1, -1, 2, 2, -1, 2, 2, 1, -2, 1, -2, 1, -2, -1]) plot_protein(p, style="basic", legend=False, show=False) - p = Protein("PHPHHPPPHHPHPHH", dim=2, model="HP") - p.set_hash([-1, -1, 2, 2, -1, 2, 2, 1, -2, 1, -2, 1, -2, -1]) - plot_protein(p, style="paper", legend_style="inner", show=False) - p = Protein("HPPHPPPHPPPPPHH", dim=2, model="HP") p.set_hash([-1, -2, 1, -2, 1, 2, 2, 1, 2, 2, -1, -1, -2, 1]) - plot_protein(p, style="paper", legend_style="inner", show=False) + plot_protein( + p, + style="paper", + legend_style="inner", + annotate_first=True, + show=False, + ) + + p = Protein("PHPPPHPPPPPHH", dim=2, model="HP") + p.set_hash([1, -2, 1, 2, 2, 1, 2, 2, -1, -1, -2, 1]) + plot_protein( + p, + style="paper", + legend_style="inner", + annotate_first=True, + show=False, + ) p = Protein("HHPHPHPPPHPHHHPPPPPPHPPPH", dim=2, model="HP") p.set_hash(