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