proteusPy.render_disulfide_schematic

  1#!/usr/bin/env python3
  2"""
  3Stand-alone program for creating 2D schematic diagrams of disulfide bonds.
  4
  5This program provides a command-line interface to the disulfide_schematic module
  6in the proteusPy package, allowing users to create publication-ready 2D diagrams
  7of disulfide bonds with various customization options.
  8
  9Author: Eric G. Suchanek, PhD
 10Last revision: 2025-03-04
 11"""
 12
 13import argparse
 14import os
 15import sys
 16from typing import Optional, Tuple
 17
 18import matplotlib.pyplot as plt
 19
 20import proteusPy as pp
 21from proteusPy.disulfide_schematic import (
 22    create_disulfide_schematic,
 23    create_disulfide_schematic_from_model,
 24)
 25
 26# Add the parent directory to the path so we can import proteusPy
 27# sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
 28
 29
 30PDB_SS = None
 31best_id = "2q7q_75D_140D"
 32worst_id = "6vxk_801B_806B"
 33
 34
 35def parse_arguments() -> argparse.Namespace:
 36    """
 37    Parse command-line arguments for the disulfide schematic renderer.
 38
 39    :return: Parsed command-line arguments
 40    :rtype: argparse.Namespace
 41    """
 42    parser = argparse.ArgumentParser(
 43        description="Create 2D schematic diagrams of disulfide bonds",
 44        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
 45    )
 46
 47    # Input source group (mutually exclusive)
 48    input_group = parser.add_mutually_exclusive_group(required=True)
 49    input_group.add_argument(
 50        "--pdb_id",
 51        type=str,
 52        help="PDB ID to load disulfides from (will use the first disulfide found)",
 53    )
 54    input_group.add_argument(
 55        "--ss_id",
 56        type=str,
 57        default=best_id,
 58        help="Specific disulfide ID to render (format: pdbid_resnum1chain_resnum2chain)",
 59    )
 60    input_group.add_argument(
 61        "--model",
 62        action="store_true",
 63        help="Create a model disulfide with specified dihedral angles",
 64    )
 65
 66    # Model disulfide parameters
 67    model_group = parser.add_argument_group("Model disulfide parameters")
 68    model_group.add_argument(
 69        "--chi1",
 70        type=float,
 71        default=-60.0,
 72        help="Chi1 dihedral angle for model disulfide",
 73    )
 74    model_group.add_argument(
 75        "--chi2",
 76        type=float,
 77        default=-60.0,
 78        help="Chi2 dihedral angle for model disulfide",
 79    )
 80    model_group.add_argument(
 81        "--chi3",
 82        type=float,
 83        default=-90.0,
 84        help="Chi3 dihedral angle for model disulfide",
 85    )
 86    model_group.add_argument(
 87        "--chi4",
 88        type=float,
 89        default=-60.0,
 90        help="Chi4 dihedral angle for model disulfide",
 91    )
 92    model_group.add_argument(
 93        "--chi5",
 94        type=float,
 95        default=-60.0,
 96        help="Chi5 dihedral angle for model disulfide",
 97    )
 98
 99    # Visualization options
100    viz_group = parser.add_argument_group("Visualization options")
101    viz_group.add_argument(
102        "--style",
103        type=str,
104        choices=["publication", "simple", "detailed"],
105        default="publication",
106        help="Visualization style",
107    )
108    viz_group.add_argument(
109        "--show_angles",
110        action="store_true",
111        default=True,
112        help="Show dihedral angles in the schematic",
113    )
114    viz_group.add_argument(
115        "--show_ca_ca_distance",
116        action="store_true",
117        default=True,
118        help="Show Cα-Cα distance in the schematic",
119    )
120    viz_group.add_argument(
121        "--show_labels",
122        action="store_true",
123        default=True,
124        help="Show atom labels in the schematic",
125    )
126    viz_group.add_argument(
127        "--hide_title",
128        action="store_true",
129        help="Hide the title in the schematic",
130    )
131    viz_group.add_argument(
132        "--dpi",
133        type=int,
134        default=300,
135        help="Resolution for raster outputs (DPI)",
136    )
137    viz_group.add_argument(
138        "--figsize",
139        type=float,
140        nargs=2,
141        default=(8, 6),
142        metavar=("WIDTH", "HEIGHT"),
143        help="Figure size in inches (width height)",
144    )
145
146    # Output options
147    output_group = parser.add_argument_group("Output options")
148    output_group.add_argument(
149        "--output_file",
150        type=str,
151        help="Path to save the output file (supports .svg, .pdf, .png)",
152    )
153    output_group.add_argument(
154        "--output_dir",
155        type=str,
156        default="schematic_outputs",
157        help="Directory to save the output file",
158    )
159    output_group.add_argument(
160        "--display",
161        action="store_true",
162        help="Display the schematic instead of saving to a file",
163    )
164
165    # Database options
166    db_group = parser.add_argument_group("Database options")
167    db_group.add_argument(
168        "--verbose",
169        action="store_true",
170        help="Show verbose output during database loading",
171    )
172
173    return parser.parse_args()
174
175
176def load_disulfide(
177    args: argparse.Namespace,
178) -> Optional[pp.DisulfideBase.Disulfide]:
179    """
180    Load a disulfide based on the provided command-line arguments.
181
182    :param args: Parsed command-line arguments
183    :type args: argparse.Namespace
184    :return: Loaded disulfide object or None if not found
185    :rtype: Optional[pp.DisulfideBase.Disulfide]
186    """
187    global PDB_SS
188
189    PDB_SS = pp.Load_PDB_SS(verbose=args.verbose, subset=False)
190
191    if args.model:
192        # Create a model disulfide with specified dihedral angles
193        model_ss = pp.DisulfideBase.Disulfide("model")
194        model_ss.build_model(args.chi1, args.chi2, args.chi3, args.chi4, args.chi5)
195        return model_ss
196
197    if args.pdb_id:
198        # Load disulfides from a specific PDB ID
199        pdb_ss = PDB_SS[args.pdb_id]
200
201        if len(pdb_ss) == 0:
202            print(f"No disulfides found for PDB ID: {args.pdb_id}")
203            return None
204        return pdb_ss[0]  # Return the first disulfide
205
206    if args.ss_id:
207        return PDB_SS[args.ss_id]
208
209    return None
210
211
212def get_output_filename(
213    args: argparse.Namespace, disulfide: pp.DisulfideBase.Disulfide
214) -> Optional[str]:
215    """
216    Determine the output filename based on command-line arguments and disulfide properties.
217
218    :param args: Parsed command-line arguments
219    :type args: argparse.Namespace
220    :param disulfide: Disulfide object to render
221    :type disulfide: pp.DisulfideBase.Disulfide
222    :return: Output filename or None if display only
223    :rtype: Optional[str]
224    """
225    if args.display:
226        return None
227
228    if args.output_file:
229        # Use the specified output file
230        if os.path.isabs(args.output_file):
231            return args.output_file
232        else:
233            return os.path.join(args.output_dir, args.output_file)
234
235    # Create a default filename based on the disulfide properties
236    if disulfide.name == "model":
237        base_name = f"model_chi1_{args.chi1}_chi2_{args.chi2}_chi3_{args.chi3}_chi4_{args.chi4}_chi5_{args.chi5}"
238    else:
239        base_name = f"{disulfide.pdb_id}_{disulfide.proximal}{disulfide.proximal_chain}_{disulfide.distal}{disulfide.distal_chain}"
240
241    filename = f"{base_name}_{args.style}.png"
242    return os.path.join(args.output_dir, filename)
243
244
245def render_disulfide_schematic(
246    disulfide: pp.DisulfideBase.Disulfide, args: argparse.Namespace
247) -> Tuple[plt.Figure, plt.Axes]:
248    """
249    Render a disulfide schematic based on the provided disulfide and arguments.
250
251    :param disulfide: Disulfide object to render
252    :type disulfide: pp.DisulfideBase.Disulfide
253    :param args: Parsed command-line arguments
254    :type args: argparse.Namespace
255    :return: Matplotlib figure and axes objects
256    :rtype: Tuple[plt.Figure, plt.Axes]
257    """
258    output_file = get_output_filename(args, disulfide)
259
260    # Create output directory if it doesn't exist and we're saving to a file
261    if output_file:
262        os.makedirs(os.path.dirname(output_file), exist_ok=True)
263
264    # Render the schematic
265    if disulfide.name == "model" and args.model:
266        # For model disulfides, use the create_disulfide_schematic_from_model function
267        fig, ax = create_disulfide_schematic_from_model(
268            chi1=args.chi1,
269            chi2=args.chi2,
270            chi3=args.chi3,
271            chi4=args.chi4,
272            chi5=args.chi5,
273            output_file=output_file,
274            show_labels=args.show_labels,
275            show_angles=args.show_angles,
276            show_ca_ca_distance=args.show_ca_ca_distance,
277            style=args.style,
278            dpi=args.dpi,
279            figsize=args.figsize,
280        )
281    else:
282        # For real disulfides, use the create_disulfide_schematic function
283        # Check if the hide_title argument was provided
284        show_title = not args.hide_title if hasattr(args, "hide_title") else True
285
286        fig, ax = create_disulfide_schematic(
287            disulfide=disulfide,
288            output_file=output_file,
289            show_labels=args.show_labels,
290            show_angles=args.show_angles,
291            show_title=show_title,
292            show_ca_ca_distance=args.show_ca_ca_distance,
293            style=args.style,
294            dpi=args.dpi,
295            figsize=args.figsize,
296        )
297
298    return fig, ax
299
300
301def main():
302    """
303    Main function for the disulfide schematic renderer.
304
305    This function parses command-line arguments, loads the specified disulfide,
306    renders the schematic, and either saves it to a file or displays it.
307    """
308    # Parse command-line arguments
309    args = parse_arguments()
310
311    # Load the disulfide
312    disulfide = load_disulfide(args)
313    if disulfide is None:
314        sys.exit(1)
315
316    # Print information about the disulfide
317    if disulfide.name == "model":
318        print(
319            f"Rendering model disulfide with angles: "
320            f"χ₁={args.chi1:.1f}°, χ₂={args.chi2:.1f}°, χ₃={args.chi3:.1f}°, "
321            f"χ₄={args.chi4:.1f}°, χ₅={args.chi5:.1f}°"
322        )
323    else:
324        print(
325            f"Rendering disulfide: {disulfide.pdb_id} "
326            f"{disulfide.proximal}{disulfide.proximal_chain}-"
327            f"{disulfide.distal}{disulfide.distal_chain}"
328        )
329        print(
330            f"Energy: {disulfide.energy:.2f} kcal/mol, "
331            f"Torsion Length: {disulfide.torsion_length:.2f}°, "
332            f"Cα Distance: {disulfide.ca_distance:.2f} Å"
333        )
334
335    # Render the schematic
336    fig, ax = render_disulfide_schematic(disulfide, args)
337
338    # Display the schematic if requested
339    if args.display:
340        print("Displaying schematic...")
341        plt.show()
342    else:
343        output_file = get_output_filename(args, disulfide)
344        print(f"Saved schematic to: {output_file}")
345
346
347if __name__ == "__main__":
348    main()