Source code for schrodinger.application.matsci.reportutils

"""
Utilities for creating MatSci reports.

Copyright Schrodinger, LLC. All rights reserved.
"""

import os
from collections import namedtuple

from matplotlib.backends import backend_pdf
from reportlab.lib.pagesizes import letter

import schrodinger.application.desmond.report_helper as rhelper
from schrodinger.application.matsci import jobutils
from schrodinger.ui.qt import filedialog
from schrodinger.ui.qt import swidgets
from schrodinger.utils import fileutils

DEFAULT_LEFT_MARGIN = 3 / 8. * rhelper.inch
DEFAULT_RIGHT_MARGIN = rhelper.inch / 4.
DEFAULT_TOP_MARGIN = 1 * rhelper.inch
DEFAULT_BOTTOM_MARGIN = 1 * rhelper.inch

EXPORT_PDF = "PDF Report"
EXPORT_IMAGES = "Plots"
EXPORT_DATA = "Data"

MATSCI_LOGO_PATH = os.path.join(fileutils.get_mmshare_scripts_dir(),
                                "event_analysis_dir", "schrodinger_ms_logo.png")

Footer = namedtuple("Footer", ["text", "height"])


[docs]class ReportInfo: """ Manages information about the report request. Contains user inputs as well as panel data. """
[docs] def __init__(self, panel, output_dir, base_name, outputs): """ Create an instance and store arguments :param `af2.App` panel: The panel for which the report is requested :param str output_dir: The directory to write the report files in :param str base_name: The base name for the exported files and folders :param list outputs: The outputs that the user has requested for the report """ self.panel = panel self.output_dir = output_dir self.base_name = base_name self.outputs = outputs self.figures = {}
[docs] def pdfRequested(self): """ Determine whether a pdf was requested for the report :rtype: bool :return: Whether a pdf was requested for the report """ return EXPORT_PDF in self.outputs
[docs] def imagesRequested(self): """ Determine whether images were requested for the report :rtype: bool :return: Whether images were requested for the report """ return EXPORT_IMAGES in self.outputs
[docs] def dataRequested(self): """ Determine whether raw data was requested for the report :rtype: bool :return: Whether raw data was requested for the report """ return EXPORT_DATA in self.outputs
[docs] def getPDFPath(self): """ Get the path to the report pdf file :rtype: str :return: path to the report pdf file """ return os.path.join(self.output_dir, self.base_name + "_report.pdf")
[docs] def getCsvFilePath(self): """ Get the path to the data csv file :rtype: str :return: path to the data csv file """ return os.path.join(self.output_dir, self.base_name + "_data.csv")
[docs] def getImagesDirPath(self): """ Get the path to the images directory :rtype: str :return: path to the images directory """ return os.path.join(self.output_dir, self.base_name + "_images")
[docs] def setFigures(self, figures): """ Set the figures for the report :param `matplotlib.figure.Figure` figures: The figures for the report """ self.figures = figures
[docs] def okToWrite(self): """ Ensure that target files and folders don't already exist, or if they do, they can be overwritten :rtype: bool :return: Whether we are clear to write """ paths_to_check = [] if self.pdfRequested(): paths_to_check.append(self.getPDFPath()) if self.imagesRequested(): paths_to_check.append(self.getImagesDirPath()) if self.dataRequested(): paths_to_check.append(self.getCsvFilePath()) existing_paths = [] for path in paths_to_check: if os.path.exists(path): existing_paths.append(os.path.basename(path)) if existing_paths: question = ("The following files or folders already exist in the" " selected directory and will be overwritten, continue?" "\n\n" + "\n".join(existing_paths)) overwrite = self.panel.question(question, button1='Yes', button2='No') if not overwrite: return False return True
[docs]class PDFBuilder: """ Contains features shared by all MatSci PDF builder classes. """
[docs] def __init__(self, report_info, left_margin=DEFAULT_LEFT_MARGIN, top_margin=DEFAULT_TOP_MARGIN, right_margin=DEFAULT_RIGHT_MARGIN, bottom_margin=DEFAULT_BOTTOM_MARGIN): """ Create a PDFBuilder instance. Store and initialize variables. :param `ReportInfo` report_info: The report information object :param int left_margin: The left margin for the pdf :param int top_margin: The top margin for the pdf :param int right_margin: The right margin for the pdf :param int bottom_margin: The bottom margin for the pdf """ self.report_info = report_info self.panel = report_info.panel self.schrodinger_temp_dir = fileutils.get_directory_path(fileutils.TEMP) self.figures = list(report_info.figures.values()) report_path = report_info.getPDFPath() self.doc = rhelper.platypus.SimpleDocTemplate(report_path, pagesize=letter) self.doc.leftMargin = left_margin self.doc.rightMargin = right_margin self.doc.topMargin = top_margin self.doc.bottomMargin = bottom_margin self.files_to_cleanup = [] self.footers = {} # Dictionary mapping page number to Footer self.elements = []
[docs] def build(self): """ Build the report pdf. Calls methods that should be defined in derived classes to add flowables and footers. """ self.addFlowables() self.addFooters() # Pass footers to the canvas to be written as the document is saved. # TODO: Find a cleaner way, such as passing a custom function to doc CanvasWithHeaderAndFooter.FOOTERS = self.footers self.doc.build(self.elements, canvasmaker=CanvasWithHeaderAndFooter) self.cleanupTempFiles()
[docs] def addFlowables(self): """ Add flowables to the pdf, such as headers, titles, paragraphs, tables and images. Should be defined in derived classes. """ raise NotImplementedError("The addFlowables method should be" " defined in all derived classes.")
[docs] def addFooters(self): """ Add footers. Should be defined in derived classes. """ raise NotImplementedError("The addFooters method should be" " defined in all derived classes.")
[docs] def writeParagraphs(self, paragraphs): """ Add the paragraphs to the pdf as flowables. :param list paragraphs: A list of strings, each of which should be a paragraph """ for paragraph in paragraphs: rhelper.pargph(self.elements, paragraph, leading=15) self.addSpacer()
[docs] def addFigure(self, figure, box_x, box_y): """ Add the figure to the pdf as a flowable. :param `matplotlib.figure.Figure` figure: The figure to add to the pdf :param box_x: The width of the box that the image should be fitted in :param box_y: The height of the box that the image should be fitted in """ assert box_x > 0 and box_y > 0 image_name = self.getTempImageFilename() figure_saved = save_figure(figure, image_name, self.panel) if not figure_saved: return size = figure.get_size_inches() resized_x, resized_y = rhelper.aspectScale(size[0], size[1], box_x, box_y) image = rhelper.platypus.Image(image_name, resized_x * rhelper.inch, resized_y * rhelper.inch) self.elements.append(image)
[docs] def getTempImageFilename(self): """ Get a temporary image file path and add it to the list of files to be cleaned up. :rtype: str :return: A temporary image file path """ temp_filename = fileutils.get_next_filename( os.path.join(self.schrodinger_temp_dir, "temp_image.png"), "") self.files_to_cleanup.append(temp_filename) return temp_filename
[docs] def cleanupTempFiles(self): """ Clean up all the temporary files """ for temp_file in self.files_to_cleanup: fileutils.force_remove(temp_file)
[docs] def getTableStyle(self, column_headers=True, row_headers=False, align_center=True): """ Create a platypus table style based on desired headers and alignment :param bool column_headers: Whether the table columns have headers :param bool row_headers: Whether the table rows have headers :param bool align_center: Whether align center should be used for all cells :rtype: list of tuples :return: A platypus table style based on desired headers and alignment """ # Format: (0, 1), (-1, -1): From column 0, row 1, to last column and row table_style = [('NOSPLIT', (0, 0), (-1, -1)), ('BOTTOMPADDING', (0, 0), (-1, -1), 1), ('TOPPADDING', (0, 0), (-1, -1), 1)] # yapf:disable if column_headers: table_style.extend([ ('BOTTOMPADDING', (0, 0), (-1, 0), 2), ('TEXTCOLOR', (0, 0), (-1, 0), rhelper.gray) ]) # yapf:disable if row_headers: table_style.append(('TEXTCOLOR', (0, 0), (0, -1), rhelper.gray)) if align_center: table_style.append(('ALIGN', (0, 0), (-1, -1), 'CENTER')) return table_style
[docs] def addSpacer(self): """ Add a space between the last and next flowables """ rhelper.add_spacer(self.elements)
[docs]class CanvasWithHeaderAndFooter(rhelper.NumberedCanvas): """ Canvas that automatically adds matsci logo to the header, and allows custom footers in addition to the page number and "Report generated" string. """ FOOTERS = {} # Dictionary mapping page number to its Footer tuple
[docs] def __init__(self, *args, **kwargs): """ Create an instance. """ super().__init__(*args, **kwargs) style = rhelper.ParaStyle style.fontName = 'Helvetica' style.fontSize = 9 style.textColor = 'black' self.footer_style = style # TODO: Passed/current margins should be used, not default ones # By extension, the font for the footer paragraph can be made customizable # They can all be set similar to Footers before building the pdf (not clean) self.doc_width = letter[0] - DEFAULT_LEFT_MARGIN - DEFAULT_RIGHT_MARGIN # Read MatSci logo and determine its position in pdf self.logo_img = rhelper.platypus.Image(MATSCI_LOGO_PATH, 1.50 * rhelper.inch, 0.43 * rhelper.inch) logo_width, logo_height = self.logo_img.wrap(0, 0) self.logo_x = letter[0] - 0.3 * rhelper.inch - logo_width self.logo_y = letter[1] - 0.5 * rhelper.inch - logo_height
[docs] def drawFixedContents(self, *args): """ Draw fixed pdf contents such as headers, footers and fixed graphics """ # Draw page number in footer super().drawFixedContents(*args) # Add extra footers if applicable page_number = self.getPageNumber() if page_number in self.FOOTERS: footer = self.FOOTERS[page_number] paragraph = rhelper.platypus.Paragraph(footer.text, self.footer_style) paragraph.wrap(self.doc_width, DEFAULT_BOTTOM_MARGIN) paragraph.drawOn(self, DEFAULT_LEFT_MARGIN, footer.height) # Draw MatSci logo self.logo_img.drawOn(self, self.logo_x, self.logo_y)
[docs]class ReportOutputsDialog(swidgets.SDialog): """ Dialog for allowing the user to specify requested report outputs. Calls the report generation method of the panel when the user accepts. """ DIR_SELECTOR_ID = "REPORT_UTILS_DIR_SELECTOR_ID"
[docs] def __init__(self, *args, default_base_name="", report_btn=True, images_btn=True, data_btn=True, **kwargs): """ Create an instance. :param bool report_btn: Whether PDF Report checkbox should be in the dialog :param bool images_btn: Whether Plots checkbox should be in the dialog :param bool data_btn: Whether Data checkbox should be in the dialog """ self.default_base_name = default_base_name self.report_btn = report_btn self.images_btn = images_btn self.data_btn = data_btn kwargs['title'] = kwargs.get('title', 'Export Options') super().__init__(*args, **kwargs)
[docs] def layOut(self): """ Lay out the widgets. """ layout = self.mylayout dator = swidgets.FileBaseNameValidator() self.base_name_le = swidgets.SLabeledEdit( "File base name:", edit_text=self.default_base_name, validator=dator, always_valid=True, layout=layout, stretch=False) self.base_name_le.setMinimumWidth(170) self.checkboxes = {} if self.report_btn: self.checkboxes[EXPORT_PDF] = swidgets.SCheckBox(EXPORT_PDF, checked=True, layout=layout) if self.images_btn: self.checkboxes[EXPORT_IMAGES] = swidgets.SCheckBox(EXPORT_IMAGES, checked=True, layout=layout) if self.data_btn: self.checkboxes[EXPORT_DATA] = swidgets.SCheckBox(EXPORT_DATA, checked=True, layout=layout)
[docs] def accept(self): """ Get user inputs and call the report generation method of the panel. """ outputs = [] for label, checkbox in self.checkboxes.items(): if checkbox.isChecked(): outputs.append(label) if not outputs: self.error("At least one option needs to be selected.") return output_dir = filedialog.get_existing_directory( self, caption="Choose Export Directory", id=self.DIR_SELECTOR_ID) if not output_dir: return report_info = ReportInfo(self.master, output_dir, self.base_name_le.text(), outputs) if not report_info.okToWrite(): return self.user_accept_function(report_info) self.info("The requested items have been exported.") return super().accept()
[docs]def save_figure(figure, file_path, panel): """ Exports the figure to a file, showing an error if it fails. :param `matplotlib.figure.Figure` figure: The figure to export :param str file_path: The path to the image file to be created :param `af2.App` panel: The panel to show any error in :rtype: bool :return: Whether the image export was successful """ # Make axis lines/ticks/text gray so that the data stands out # Need to revert to original color after exporting image # for axes in figure.axes: # rhelper.change_plot_colors(axes) try: figure.savefig(file_path, bbox_inches='tight', dpi=300) except (IOError, OSError) as err: panel.error(f"Error saving {file_path}: {err}") return False return True
[docs]def save_images(report_info): """ Save all images in the report information object to the output directory :param `ReportInfo` report_info: The report information object containing the figures """ image_dir = report_info.getImagesDirPath() fileutils.force_rmtree(image_dir) fileutils.mkdir_p(image_dir) for title, figure in report_info.figures.items(): file_path = os.path.join(image_dir, title + ".png") save_figure(figure, file_path, report_info.panel)
[docs]def sub_script(text, size=10): """ :type text: str or int or float :param text: The text to write as a subscript :param int size: The font size of the subscript :rtype: str :return: The text formatted to appear as a subscript in the report """ return f"<sub><font size={size}>{text}</font></sub>"
[docs]def super_script(text, size=10): """ :type text: str or int or float :param text: The text to write as a superscript :param int size: The font size of the superscript :rtype: str :return: The text formatted to appear as a superscript in the report """ return f"<super><font size={size}>{text}</font></super>"
[docs]def save_matplotlib_figure(figures, pdf_name): """ Save or Append the matplotlib figure to a pdf :type figures: list :param figures: list of matplotlib figure to be added to the pdf :type pdf_name: str :param pdf_name: filename for the pdf file """ with backend_pdf.PdfPages(pdf_name) as pdf_pages: for figure in figures: pdf_pages.savefig(figure) jobutils.add_outfile_to_backend(pdf_name)