<commit_before>#! /usr/bin/env python
"""Make static images of lyman results using PySurfer."""
import os
import os.path as op
import sys
import tempfile
import argparse
from textwrap import dedent
from time import sleep

import numpy as np
from scipy import stats
import nibabel as nib
import matplotlib.image as mplimg
from mayavi import mlab
from surfer import Brain, io

import moss
import lyman


def main(arglist):

    # Parse the command line
    args = parse_args(arglist)

    # Configure mayavi
    mlab.options.offscreen = args.nowindow

    # Load the lyman data
    subjects = lyman.determine_subjects(args.subjects)
    project = lyman.gather_project_info()
    exp = lyman.gather_experiment_info(args.experiment, args.altmodel)
    contrasts = exp["contrast_names"]
    z_thresh = exp["cluster_zthresh"]

    # Get the full correct name for the experiment
    if args.experiment is None:
        exp_name = project["default_exp"]
    else:
        exp_name = args.experiment
    exp_base = exp_name
    if args.altmodel is not None:
        exp_name = "-".join([exp_base, args.altmodel])

    # Group-level
    # ===========

    if args.level == "group":
        temp_base = op.join(project["analysis_dir"], exp_name, args.output,
                            args.regspace, "{contrast}")
        if args.regspace == "fsaverage":
            sig_thresh = -np.log10(stats.norm.sf(z_thresh))
            sig_thresh = np.round(sig_thresh) * 10
            corr_sign = exp["surf_corr_sign"]
            sig_name = "cache.th%d.%s.sig.masked.mgh" % (sig_thresh, corr_sign)
            stat_temp = op.join(temp_base, "{hemi}/osgm", sig_name)
            mask_temp = op.join(temp_base, "{hemi}/mask.mgh")
            png_temp = op.join(temp_base, "{hemi}/osgm/zstat_threshold.png")
        else:
            stat_temp = op.join(temp_base, "{hemi}.zstat1_threshold.mgz")
            mask_temp = op.join(temp_base, "{hemi}.group_mask.mgz")
            png_temp = op.join(temp_base, "zstat1_threshold_surf.png")
            corr_sign = "pos"

        contrast_loop("fsaverage", contrasts, stat_temp, mask_temp, png_temp,
                      args, z_thresh, corr_sign)

    # Subject-level
    # =============

    elif args.level == "subject":
        temp_base = op.join(project["analysis_dir"], exp_name, "{subj}",
                            "ffx", args.regspace, "smoothed/{contrast}")
        mask_temp = op.join(temp_base, "{hemi}.mask.mgz")
        stat_temp = op.join(temp_base, "{hemi}.zstat1.mgz")
        png_temp = op.join(temp_base, "zstat1_surf.png")

        for subj in subjects:
            contrast_loop(subj, contrasts, stat_temp, mask_temp, png_temp,
                          args, 1.96, "abs")


def contrast_loop(subj, contrasts, stat_temp, mask_temp, png_temp,
                  args, z_thresh, sign):
    """Iterate over contrasts and make surface images."""
    config_opts = dict(background="white", width=600, height=370)

    for contrast in contrasts:

        # Calculate where the overlay should saturate
        z_max = calculate_sat_point(stat_temp, contrast, subj)
        png_panes = []
        for hemi in ["lh", "rh"]:

            # Initialize the brain object
            b_subj = subj if args.regspace == "epi" else "fsaverage"
            b = Brain(b_subj, hemi, args.geometry,
                      config_opts=config_opts)

            # Plot the mask
            mask_file = mask_temp.format(contrast=contrast,
                                         hemi=hemi, subj=subj)
            add_mask_overlay(b, mask_file)

            # Plot the overlay
            stat_file = stat_temp.format(contrast=contrast,
                                         hemi=hemi, subj=subj)
            add_stat_overlay(b, stat_file, z_thresh, z_max, sign,
                             sig_to_z=args.regspace == "fsaverage")
            png_panes.append(save_view_panes(b, sign, hemi))

            # Maybe close the current figure
            if args.close:
                b.close()

        # Stitch the hemisphere pngs together and save
        full_png = np.concatenate(png_panes, axis=1)
        png_file = png_temp.format(contrast=contrast, hemi=hemi, subj=subj)
        mplimg.imsave(png_file, full_png)


def calculate_sat_point(template, contrast, subj=None):
    """Calculate the point at which the colormap should saturate."""
    data = []
    for hemi in ["lh", "rh"]:
        hemi_file = template.format(contrast=contrast, subj=subj, hemi=hemi)
        hemi_data = nib.load(hemi_file).get_data()
        data.append(hemi_data)
    data = np.concatenate(data)
    return max(3.71, moss.percentiles(data, 98))


def add_mask_overlay(b, mask_file):
    """Gray-out vertices outside of the common-space mask."""
    mask_data = nib.load(mask_file).get_data()

    # Plot the mask
    mask_data = np.logical_not(mask_data.astype(bool)).squeeze()
    if mask_data.any():
        b.add_data(mask_data, min=0, max=10, thresh=.5,
                   colormap="bone", alpha=.6, colorbar=False)


def add_stat_overlay(b, stat_file, thresh, max, sign, sig_to_z=False):
    """Plot a surface-encoded statistical overlay."""
    stat_data = nib.load(stat_file).get_data()

    # Possibly convert -log10(p) images to z stats
    if sig_to_z:
        stat_sign = np.sign(stat_data)
        p_data = 10 ** -np.abs(stat_data)
        z_data = stats.norm.ppf(p_data)
        z_data[np.sign(z_data) != stat_sign] *= -1
        stat_data = z_data

    # Plot the statistical data
    plot_stat = ((sign == "pos" and (stat_data > thresh).any())
                 or (sign == "neg" and (stat_data < -thresh).any())
                 or (sign == "abs" and (np.abs(stat_data) > thresh).any()))
    if plot_stat:
        stat_data = stat_data.squeeze()
        b.add_overlay(stat_data, min=thresh, max=max, sign=sign, name="stat")


def remove_stat_overlay(b):
    """Remove any overlays from the brain."""
    for name, overlay in b.overlays.items():
        overlay.remove()


def save_view_panes(b, sign, hemi):
    """Save lat, med and ven views, return a stacked image array."""
    views = ["lat", "med", "ven"]
    image_panes = []
    for view in views:

        # Handle the colorbar
        if "stat" in b.overlays:
            show_pos = hemi == "rh" and view == "ven"
            show_neg = hemi == "lh" and view == "ven"
            if sign in ("pos", "abs"):
                b.overlays["stat"].pos_bar.visible = show_pos
            if sign in ("neg", "abs"):
                b.overlays["stat"].neg_bar.visible = show_neg

        # Set the view and screenshot
        b.show_view(view, distance=330)
        pause(0.2)  # pause for show_view
        image_panes.append(b.screenshot())

    # Stitch the images together and return the array
    full_image = np.concatenate(image_panes, axis=0)

    return full_image


def parse_args(arglist):
<commit_msg>sleep to allow show_view to catch up<commit_after>#! /usr/bin/env python
"""Make static images of lyman results using PySurfer."""
import os
import os.path as op
import sys
import tempfile
import argparse
from textwrap import dedent
from time import sleep

import numpy as np
from scipy import stats
import nibabel as nib
import matplotlib.image as mplimg
from mayavi import mlab
from surfer import Brain, io

import moss
import lyman


def main(arglist):

    # Parse the command line
    args = parse_args(arglist)

    # Configure mayavi
    mlab.options.offscreen = args.nowindow

    # Load the lyman data
    subjects = lyman.determine_subjects(args.subjects)
    project = lyman.gather_project_info()
    exp = lyman.gather_experiment_info(args.experiment, args.altmodel)
    contrasts = exp["contrast_names"]
    z_thresh = exp["cluster_zthresh"]

    # Get the full correct name for the experiment
    if args.experiment is None:
        exp_name = project["default_exp"]
    else:
        exp_name = args.experiment
    exp_base = exp_name
    if args.altmodel is not None:
        exp_name = "-".join([exp_base, args.altmodel])

    # Group-level
    # ===========

    if args.level == "group":
        temp_base = op.join(project["analysis_dir"], exp_name, args.output,
                            args.regspace, "{contrast}")
        if args.regspace == "fsaverage":
            sig_thresh = -np.log10(stats.norm.sf(z_thresh))
            sig_thresh = np.round(sig_thresh) * 10
            corr_sign = exp["surf_corr_sign"]
            sig_name = "cache.th%d.%s.sig.masked.mgh" % (sig_thresh, corr_sign)
            stat_temp = op.join(temp_base, "{hemi}/osgm", sig_name)
            mask_temp = op.join(temp_base, "{hemi}/mask.mgh")
            png_temp = op.join(temp_base, "{hemi}/osgm/zstat_threshold.png")
        else:
            stat_temp = op.join(temp_base, "{hemi}.zstat1_threshold.mgz")
            mask_temp = op.join(temp_base, "{hemi}.group_mask.mgz")
            png_temp = op.join(temp_base, "zstat1_threshold_surf.png")
            corr_sign = "pos"

        contrast_loop("fsaverage", contrasts, stat_temp, mask_temp, png_temp,
                      args, z_thresh, corr_sign)

    # Subject-level
    # =============

    elif args.level == "subject":
        temp_base = op.join(project["analysis_dir"], exp_name, "{subj}",
                            "ffx", args.regspace, "smoothed/{contrast}")
        mask_temp = op.join(temp_base, "{hemi}.mask.mgz")
        stat_temp = op.join(temp_base, "{hemi}.zstat1.mgz")
        png_temp = op.join(temp_base, "zstat1_surf.png")

        for subj in subjects:
            contrast_loop(subj, contrasts, stat_temp, mask_temp, png_temp,
                          args, 1.96, "abs")


def contrast_loop(subj, contrasts, stat_temp, mask_temp, png_temp,
                  args, z_thresh, sign):
    """Iterate over contrasts and make surface images."""
    config_opts = dict(background="white", width=600, height=370)

    for contrast in contrasts:

        # Calculate where the overlay should saturate
        z_max = calculate_sat_point(stat_temp, contrast, subj)
        png_panes = []
        for hemi in ["lh", "rh"]:

            # Initialize the brain object
            b_subj = subj if args.regspace == "epi" else "fsaverage"
            b = Brain(b_subj, hemi, args.geometry,
                      config_opts=config_opts)

            # Plot the mask
            mask_file = mask_temp.format(contrast=contrast,
                                         hemi=hemi, subj=subj)
            add_mask_overlay(b, mask_file)

            # Plot the overlay
            stat_file = stat_temp.format(contrast=contrast,
                                         hemi=hemi, subj=subj)
            add_stat_overlay(b, stat_file, z_thresh, z_max, sign,
                             sig_to_z=args.regspace == "fsaverage")
            png_panes.append(save_view_panes(b, sign, hemi))

            # Maybe close the current figure
            if args.close:
                b.close()

        # Stitch the hemisphere pngs together and save
        full_png = np.concatenate(png_panes, axis=1)
        png_file = png_temp.format(contrast=contrast, hemi=hemi, subj=subj)
        mplimg.imsave(png_file, full_png)


def calculate_sat_point(template, contrast, subj=None):
    """Calculate the point at which the colormap should saturate."""
    data = []
    for hemi in ["lh", "rh"]:
        hemi_file = template.format(contrast=contrast, subj=subj, hemi=hemi)
        hemi_data = nib.load(hemi_file).get_data()
        data.append(hemi_data)
    data = np.concatenate(data)
    return max(3.71, moss.percentiles(data, 98))


def add_mask_overlay(b, mask_file):
    """Gray-out vertices outside of the common-space mask."""
    mask_data = nib.load(mask_file).get_data()

    # Plot the mask
    mask_data = np.logical_not(mask_data.astype(bool)).squeeze()
    if mask_data.any():
        b.add_data(mask_data, min=0, max=10, thresh=.5,
                   colormap="bone", alpha=.6, colorbar=False)


def add_stat_overlay(b, stat_file, thresh, max, sign, sig_to_z=False):
    """Plot a surface-encoded statistical overlay."""
    stat_data = nib.load(stat_file).get_data()

    # Possibly convert -log10(p) images to z stats
    if sig_to_z:
        stat_sign = np.sign(stat_data)
        p_data = 10 ** -np.abs(stat_data)
        z_data = stats.norm.ppf(p_data)
        z_data[np.sign(z_data) != stat_sign] *= -1
        stat_data = z_data

    # Plot the statistical data
    plot_stat = ((sign == "pos" and (stat_data > thresh).any())
                 or (sign == "neg" and (stat_data < -thresh).any())
                 or (sign == "abs" and (np.abs(stat_data) > thresh).any()))
    if plot_stat:
        stat_data = stat_data.squeeze()
        b.add_overlay(stat_data, min=thresh, max=max, sign=sign, name="stat")


def remove_stat_overlay(b):
    """Remove any overlays from the brain."""
    for name, overlay in b.overlays.items():
        overlay.remove()


def save_view_panes(b, sign, hemi):
    """Save lat, med and ven views, return a stacked image array."""
    views = ["lat", "med", "ven"]
    image_panes = []
    for view in views:

        # Handle the colorbar
        if "stat" in b.overlays:
            show_pos = hemi == "rh" and view == "ven"
            show_neg = hemi == "lh" and view == "ven"
            if sign in ("pos", "abs"):
                b.overlays["stat"].pos_bar.visible = show_pos
            if sign in ("neg", "abs"):
                b.overlays["stat"].neg_bar.visible = show_neg

        # Set the view and screenshot
        b.show_view(view, distance=330)
        sleep(0.2)  # pause for show_view
        image_panes.append(b.screenshot())

    # Stitch the images together and return the array
    full_image = np.concatenate(image_panes, axis=0)

    return full_image


def parse_args(arglist):

{'fraction_non_alphanumeric': 0.11810738352293312, 'fraction_numerical': 0.005651764364901094, 'mean_word_length': 2.4358974358974357, 'pattern_counts': {'":': 6, '<': 5, '<?xml version=': 0, '>': 7, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 8, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3839402', 'n_tokens_mistral': 4655, 'n_tokens_neox': 4422, 'n_words': 1062}