text:
#!/usr/bin/env python3

import csv
import os
import subprocess

# Import libaries we need
from matplotlib.gridspec import GridSpec
from scipy.interpolate import interp1d, RectBivariateSpline
from scipy.linalg import solve
from sklearn.model_selection import ParameterGrid

# Import user functions
from mad_trax import *

# Exceptions
class UnstableSolutionException(Exception):
    pass

class MissingOption(Exception):
    pass

# Unit conversions
def yr2sec(time):
    """Converts time from years to seconds."""
    return time * 60.0 * 60.0 * 24.0 * 365.25

def myr2sec(time):
    """Converts time from million years to seconds."""
    return yr2sec(time) * 1.0e6

def kilo2base(value):
    """Converts value from kilo-units to the base unit."""
    return value * 1000.0

def milli2base(value):
    """Converts value from milli-units to the base unit."""
    return value / 1.0e3

def micro2base(value):
    """Converts value from micro-units to the base unit."""
    return value / 1.0e6

def mmyr2ms(rate):
    """Converts rate from mm/yr to m/s."""
    return milli2base(rate) / yr2sec(1)

def echo_model_info(
    dx,
    nt,
    dt,
    t_total,
    implicit,
    vx,
    k_crust,
    rho_crust,
    cp_crust,
    k_mantle,
    rho_mantle,
    cp_mantle,
    k_a,
    erotype,
    cond_crit=0.5,
    adv_crit=0.5,
):
    print("")
    print("--- General model information ---")
    print("")
    print("- Node spacing: {0} m".format(dx))
    print("- Total simulation time: {0:.1f} million years".format(t_total / myr2sec(1)))
    print("- Time steps: {0} @ {1:.1f} years each".format(nt, dt / yr2sec(1)))

    if implicit:
        print("- Solution type: Implicit")
    else:
        print("- Solution type: Explicit")

    # Check stability conditions
    if not implicit:
        kappa_crust = k_crust / (rho_crust * cp_crust)
        kappa_mantle = k_mantle / (rho_mantle * cp_mantle)
        kappa_a = k_a / (rho_mantle * cp_mantle)
        kappa = max(kappa_crust, kappa_mantle, kappa_a)
        cond_stab = kappa * dt / dx**2
        print(
            "- Conductive stability: {0} ({1:.3f} < {2:.4f})".format(
                (cond_stab < cond_crit), cond_stab, cond_crit
            )
        )
        if cond_stab >= cond_crit:
            raise UnstableSolutionException(
                "Heat conduction solution unstable. Decrease nx or dt."
            )

        adv_stab = vx * dt / dx
        print(
            "- Advective stability: {0} ({1:.3f} < {2:.4f})".format(
                (adv_stab < adv_crit), adv_stab, adv_crit
            )
        )
        if adv_stab >= adv_crit:
            raise UnstableSolutionException(
                "Heat advection solution unstable. Decrease nx, dt, or vx (change in Moho over model time)."
            )

    # Output erosion model
    ero_models = {1: "Constant", 2: "Step-function", 3: "Exponential decay", 4: "Thrust sheet emplacement/erosion"}
    print("- Erosion model: {0}".format(ero_models[erotype]))

# Mantle adiabat from Turcotte and Schubert (eqn 4.254)
def adiabat(alphav, temp, cp):
    """Calculates a mantle adiabat in degress / m."""
    grav = 9.81
    return alphav * grav * temp / cp

# Conductive steady-state heat transfer
def temp_ss_implicit(nx, dx, temp_surf, temp_base, vx, rho, cp, k, heat_prod):
    """Calculates a steady-state thermal solution."""
    # Create the empty (zero) coefficient and right hand side arrays
    a_matrix = np.zeros((nx, nx))  # 2-dimensional array, ny rows, ny columns
    b = np.zeros(nx)

    # Set B.C. values in the coefficient array and in the r.h.s. array
    a_matrix[0, 0] = 1
    b[0] = temp_surf
    a_matrix[nx - 1, nx - 1] = 1
    b[nx - 1] = temp_base

    # Matrix loop
    for ix in range(1, nx - 1):
        a_matrix[ix, ix - 1] = (-(rho[ix - 1] * cp[ix - 1] * -vx) / (2 * dx)) - k[
            ix - 1
        ] / dx**2
        a_matrix[ix, ix] = k[ix] / dx**2 + k[ix - 1] / dx**2
        a_matrix[ix, ix + 1] = (rho[ix + 1] * cp[ix + 1] * -vx) / (2 * dx) - k[
            ix
        ] / dx**2
        b[ix] = heat_prod[ix]

    temp = solve(a_matrix, b)
    return temp

def update_materials(
    x,
    xstag,
    moho_depth,
    rho_crust,
    rho_mantle,
    rho,
    cp_crust,
    cp_mantle,
    cp,
    k_crust,
    k_mantle,
    k,
    heat_prod_crust,
    heat_prod_mantle,
    heat_prod,
    temp_adiabat,
    temp_prev,
    k_a,
    removal_fraction,
):
    """Updates arrays of material properties."""
    rho[:] = rho_crust
    rho[x > moho_depth] = rho_mantle
    cp[:] = cp_crust
    cp[x > moho_depth] = cp_mantle
    k[:] = k_crust
    k[xstag > moho_depth] = k_mantle

    interp_temp_prev = interp1d(x, temp_prev)
    temp_stag = interp_temp_prev(xstag)
    k[temp_stag >= temp_adiabat] = k_a
    if removal_fraction > 0.0:
        lab_depth = xstag[temp_stag >= temp_adiabat].min()
    else:
        lab_depth = x.max()

    heat_prod[:] = heat_prod_crust
    heat_prod[x > moho_depth] = heat_prod_mantle
    return rho, cp, k, heat_prod, lab_depth

def init_erotype4(params, temp_init, x, xstag, temp_prev, moho_depth):
    """Defines temperatures and material properties for erotype 4 (thrust sheet emplacement)."""

    # Find index where depth reaches or exceeds thrust sheet thickness
    ref_index = np.min(np.where(x >= kilo2base(params["erotype_opt1"])))

    # Reassign temperatures
    for ix in range(params["nx"]):
        if ix >= ref_index:
            temp_prev[ix] = temp_init[ix-ref_index]

    # Modify moho_depth, material property arrays
    moho_depth += kilo2base(params["erotype_opt1"])
    rho = np.ones(len(x)) * params["rho_crust"]
    rho[x > moho_depth] = params["rho_mantle"]
    cp = np.ones(len(x)) * params["cp_crust"]
    cp[x > moho_depth] = params["cp_mantle"]
    k = np.ones(len(xstag)) * params["k_crust"]
    k[xstag > moho_depth] = params["k_mantle"]
    heat_prod = np.ones(len(x)) * micro2base(params["heat_prod_crust"])
    heat_prod[x > moho_depth] = micro2base(params["heat_prod_mantle"])
    alphav = np.ones(len(x)) * params["alphav_crust"]
    alphav[x > moho_depth] = params["alphav_mantle"]

    return temp_prev, moho_depth, rho, cp, k, heat_prod, alphav

def temp_transient_explicit(
    temp_prev, temp_new, temp_surf, temp_base, nx, dx, vx, dt, rho, cp, k, heat_prod
):
    """Updates a transient thermal solution."""
    # Set boundary conditions
    temp_new[0] = temp_surf
    temp_new[nx - 1] = temp_base

    # Calculate internal grid point temperatures
    for ix in range(1, nx - 1):
        temp_new[ix] = (
            (1 / (rho[ix] * cp[ix]))
            * (
                k[ix] * (temp_prev[ix + 1] - temp_prev[ix])
                - k[ix - 1] * (temp_prev[ix] - temp_prev[ix - 1])
            )
            / dx**2
            + heat_prod[ix] / (rho[ix] * cp[ix])
            + vx * (temp_prev[ix + 1] - temp_prev[ix - 1]) / (2 * dx)
        ) * dt + temp_prev[ix]

    return temp_new

# Conductive steady-state heat transfer
def temp_transient_implicit(
    nx, dx, dt, temp_prev, temp_surf, temp_base, vx, rho, cp, k, heat_prod
):
    """Calculates a steady-state thermal solution."""
    # Create the empty (zero) coefficient and right hand side arrays
    a_matrix = np.zeros((nx, nx))  # 2-dimensional array, ny rows, ny columns
    b = np.zeros(nx)

    # Set B.C. values in the coefficient array and in the r.h.s. array
    a_matrix[0, 0] = 1
    b[0] = temp_surf
    a_matrix[nx - 1, nx - 1] = 1
    b[nx - 1] = temp_base

    # Matrix loop
    for ix in range(1, nx - 1):
        a_matrix[ix, ix - 1] = (
            -(rho[ix - 1] * cp[ix - 1] * -vx) / (2 * dx) - k[ix - 1] / dx**2
        )
        a_matrix[ix, ix] = (
            (rho[ix] * cp[ix]) / dt + k[ix] / dx**2 + k[ix - 1] / dx**2
        )
        a_matrix[ix, ix + 1] = (rho[ix + 1] * cp[ix + 1] * -vx) / (2 * dx) - k[
            ix
        ] / dx**2
        b[ix] = heat_prod[ix] + ((rho[ix] * cp[ix]) / dt) * temp_prev[ix]

    temp = solve(a_matrix, b)
    return temp

def he_ages(
    file,
    ap_rad=45.0,
    ap_uranium=10.0,
    ap_thorium=40.0,
    zr_rad=60.0,
    zr_uranium=100.0,
    zr_thorium=40.0,
):
    """Calculates (U-Th)/He ages."""

    command = (
        "../bin/RDAAM_He "
        + file
        + " "
        + str(ap_rad)
        + " "
        + str(ap_uranium)
        + " "
        + str(ap_thorium)
        + " "
        + str(zr_rad)
        + " "
        + str(zr_uranium)
        + " "
        + str(zr_thorium)
    )
    p = subprocess.Popen(
        command, shell=True, stdout=subprocess.PIPE, stderr=subprocess.STDOUT
    )

    stdout = p.stdout.readlines()

    ahe_age = stdout[0].split()[3][:-1].decode("UTF-8")
    corr_ahe_age = stdout[0].split()[7].decode("UTF-8")
    zhe_age = stdout[1].split()[3][:-1].decode("UTF-8")
    corr_zhe_age = stdout[1].split()[7].decode("UTF-8")

    retval = p.wait()
    return ahe_age, corr_ahe_age, zhe_age, corr_zhe_age

def ft_ages(file):
    """Calculates AFT ages."""

    command = "../bin/ketch_aft " + file
    p = subprocess.Popen(
        command, shell=True, stdout=subprocess.PIPE, stderr=subprocess.STDOUT
    )

    stdout = p.stdout.readlines()
    aft_age = stdout[0].split()[4][:-1].decode("UTF-8")
    mean_ft_length = stdout[0].split()[9][:-1].decode("UTF-8")

    retval = p.wait()
    return aft_age, mean_ft_length

def calculate_erosion_rate(
    t_total, current_time, magnitude, erotype, erotype_opt1, erotype_opt2
):
    """Defines the way in which erosion should be applied."""

    # Constant erosion rate
    if erotype == 1:
        vx = magnitude / t_total

    # Constant erosion rate with a step-function change at a specified time
    elif erotype == 2:
        init_rate = mmyr2ms(erotype_opt1)
        rate_change_time = myr2sec(erotype_opt2)
        remaining_magnitude = magnitude - (init_rate * rate_change_time)
        # First stage of erosion
        if current_time < rate_change_time:
            vx = init_rate
        # Second stage of erosion
        else:
            vx = remaining_magnitude / (t_total - rate_change_time)

    # Exponential erosion rate decay with a set characteristic time
    elif erotype == 3:
        decay_time = myr2sec(erotype_opt1)
        # Calculate max erosion rate for exponential
        max_rate = magnitude / (
            decay_time * (np.exp(0.0 / decay_time) - np.exp(-t_total / decay_time))
        )
        vx = max_rate * np.exp(-current_time / decay_time)

    # Emplacement and erosional removal of a thrust sheet
    elif erotype == 4:
        vx = kilo2base(erotype_opt1)/t_total

    # Catch bad cases
    else:
        raise MissingOption("Bad erosion type. Type should be 1, 2, 3, or 4.")

    return vx

def calculate_pressure(density, dx, g=9.81):
    """Calculates lithostatic pressure"""
    pressure = np.zeros(len(density))
    for i in range(1, len(density)):
        pressure[i] = pressure[i - 1] + density[i] * g * dx

    return pressure

def calculate_crust_solidus(composition, crustal_pressure):
    """Reads in data from MELTS for different compositions and returns a solidus"""

    # Composition options
    compositions = {
        "wet_felsic": "wetFelsic4.csv",
        "wet_intermediate": "wetImed3.csv",
        "wet_basalt": "wetBasalt2.csv",
        "dry_felsic": "dryFelsic.csv",
        "dry_basalt": "dryBasalt.csv",
    }

    # Read composition data file
    fp = "melts_data/" + compositions[composition]
    comp_data = np.genfromtxt(fp, delimiter=",")

    # Create interpolation function for composition
    crust_interp = RectBivariateSpline(
        comp_data[2:, 0], comp_data[0, 1:], comp_data[2:, 1:], kx=1, ky=1
    )

    # Creating Pressure vs Melts fraction grid which gives the values for temperatures
    Tn = np.linspace(0, 1, 121)  # Last number defines the number of melt fraction steps
    Tn[-1] = 0.99999999  # Highest temperature value
    interp_list = []  # lists of melt fractions at interpolated pressure ranges

    for i in range(len(comp_data[0, 1:])):
        interp_list.append(
            np.interp(Tn, comp_data[2:, i + 1], comp_data[2:, 0])
        )  # x(melt), y(T)

    # Creating 2D array of Pressure (x axis) vs melt fraction (y axis)
    interp_list = np.transpose(np.array(interp_list))

    # Interpolating the melt fraction vs pressure data (gives temperature)
    interp_temp = RectBivariateSpline(Tn, comp_data[0, 1:], interp_list, kx=1, ky=1)

    # Melting curve plots for the materials
    # melt_fractions = np.array([0.05, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1.0])
    melt_fractions = np.array([0.05])

    for i in range(len(melt_fractions)):
        crust_solidus = interp_temp(melt_fractions[i], crustal_pressure)

    return crust_solidus[0, :]

def calculate_mantle_solidus(pressure, xoh=0.0):
    """Calculates the solidus for the mantle.
    Uses the equation from Hirschmann, 2000 as modified by Sarafian et al. (2017)."""

    # Convert ppm to parts
    xoh *= 1.0e-6

    # Hirschmann constants
    a = -5.104
    b = 132.899
    c = 1120.661

    # Hirschmann solidus
    solidus = a * pressure**2 + b * pressure + c

    # Sarafian modifications
    gas_constant = 8.314
    silicate_mols = 59.0
    entropy = 0.4
    solidus = solidus / (
        1 - (gas_constant / (silicate_mols * entropy)) * np.log(1 - xoh)
    )

    return solidus

def calculate_misfit(predicted_ages, measured_ages, measured_stdev, num_params, type):
    """
    Calculates misfit value between measured and predicted thermochronometer ages

    type 1 = Braun et al. (2012) equation 8
    type 2 = Braun et al. (2012) equation 9
    type 3 = Braun et al. (2012) equation 10
    """

    if type == 1:
        misfit = np.sqrt(
            ((predicted_ages - measured_ages) ** 2 / measured_stdev**2).sum()
        ) / len(predicted_ages)

    if type == 2:
        misfit = ((predicted_ages - measured_ages) ** 2 / measured_stdev**2).sum() / (
            len(predicted_ages) - num_params - 1
        )

    if type == 3:
        misfit = (((predicted_ages - measured_ages) / measured_stdev) ** 2).sum()

    return misfit

def prep_model(params):
    """Prepares models to be run as single models or in batch mode"""

    batch_keys = [
        "max_depth",
        "nx",
        "temp_surf",
        "temp_base",
        "t_total",
        "dt",
        "vx_init",
        "init_moho_depth",
        "final_moho_depth",
        "removal_fraction",
        "crustal_flux",
        "erotype",
        "erotype_opt1",
        "erotype_opt2",
        "mantle_adiabat",
        "rho_crust",
        "cp_crust",
        "k_crust",
        "heat_prod_crust",
        "alphav_crust",
        "rho_mantle",
        "cp_mantle",
        "k_mantle",
        "heat_prod_mantle",
        "alphav_mantle",
        "rho_a",
        "k_a",
        "ap_rad",
        "ap_uranium",
        "ap_thorium",
        "zr_rad",
        "zr_uranium",
        "zr_thorium",
        "pad_thist",
        "pad_time",
    ]

    # Create empty dictionary for batch model parameters, if any
    batch_params = {}

    # Check that prep_model was called from the command line
    if params["cmd_line_call"]:
        # We know all batch model values are lists, check their lengths
        # If all are 1 then run in single-model mode
        params["batch_mode"] = False

        for key in batch_keys:
            if len(params[key]) != 1:
                params["batch_mode"] = True
            batch_params[key] = params[key]

        # Now we see what to do for running the model(s)
        if not params["batch_mode"]:
            # Convert list values, run single model
            for key in batch_keys:
                params[key] = params[key][0]
            run_model(params)
        else:
            # Run in batch mode
            batch_run(params, batch_params)

    else:
        # If called as a function, check for lists and their lengths
        params["batch_mode"] = False

        for key in batch_keys:
            if isinstance(params[key], list):
                if len(params[key]) != 1:
                    params["batch_mode"] = True
                batch_params[key] = params[key]

        # Now we see what to do for running the model(s)
        if not params["batch_mode"]:
            # Convert list values, run single model
            run_model(params)
        else:
            # Run in batch mode
            batch_run(params, batch_params)

def batch_run(params, batch_params):
    """Runs TC1D in batch mode"""
    param_list = list(ParameterGrid(batch_params))

    print("--- Starting batch processor for {0} models ---\n".format(len(param_list)))

    # Check number of past models and write header as needed
    # Define output file
    outfile = "TC1D_batch_log.csv"

    # Open file for writing
    model_count = 0
    try:
        with open(outfile) as f:
            write_header = False
            infile = f.readlines()
            if len(infile) < 1:
                write_header = True
            else:
                model_count = len(infile) - 1
    except IOError:
        write_header = True

    success = 0
    failed = 0

    for i in range(len(param_list)):
        model_count += 1
        model_id = "M{0}".format(str(model_count).zfill(4))
        model = param_list[i]
        print("Iteration {0}...".format(i + 1), end="", flush=True)
        # Update model parameters
        for key in batch_params:
            params[key] = model[key]

        # Open file for writing
        with open(outfile, "a+") as f:
            if write_header:
                f.write(
                    "Model ID,Simulation time (Myr),Time step (yr),Model thickness (km),Node points,"
                    "Surface temperature (C),Basal temperature (C),Mantle adiabat,"
                    "Crustal density (kg m^-3),Mantle removal fraction,"
                    "Erosion model type,Erosion model option 1,"
                    "Erosion model option 2,Initial Moho depth (km),Initial Moho temperature (C),"
                    "Initial surface heat flow (mW m^-2),Initial surface elevation (km),"
                    "Final Moho depth (km),Final Moho temperature (C),Final surface heat flow (mW m^-2),"
                    "Final surface elevation (km),Apatite grain radius (um),Apatite U concentration (ppm),"
                    "Apatite Th concentration (ppm),Zircon grain radius (um),Zircon U concentration (ppm),"
                    "Zircon Th concentration (ppm),Predicted apatite (U-Th)/He age (Ma),"
                    "Predicted apatite (U-Th)/He closure temperature (C),Measured apatite (U-Th)/He age (Ma),"
                    "Measured apatite (U-Th)/He standard deviation (Ma),Predicted apatite fission-track age (Ma),"
                    "Predicted apatite fission-track closure temperature (C),Measured apatite fission-track age (Ma),"
                    "Measured apatite fission-track standard deviation (Ma),Predicted zircon (U-Th)/He age (Ma),"
                    "Predicted zircon (U-Th)/He closure temperature (C),Measured zircon (U-Th)/He age (Ma),"
                    "Measured zircon (U-Th)/He standard deviation (Ma),Misfit,Misfit type,Number of ages for misfit\n"
                )
                write_header = False
            f.write("{0},".format(model_id))
        params["model_id"] = model_id

        try:
            run_model(params)
            print("{0}".format("Complete"))
            success += 1
        except:
            print("{0}".format("FAILED!"))
            with open(outfile, "a+") as f:
                f.write(
                    "{0:.4f},{1:.4f},{2:.4f},{3},{4:.4f},{5:.4},{6},{7:.4f},"
                    "{8:.4f},{9},{10:.4f},{11:.4f},{12:.4f},,,,{13:.4f},"
                    ",,,{14:.4f},{15:.4f},{16:.4f},{17:.4f},{18:.4f},"
                    "{19:.4f},,,,,,,,,,,,,,,"
                    "\n".format(
                        params["t_total"],
                        params["dt"],
                        params["max_depth"],
                        params["nx"],
                        params["temp_surf"],
                        params["temp_base"],
                        params["mantle_adiabat"],
                        params["rho_crust"],
                        params["removal_fraction"],
                        params["erotype"],
                        params["erotype_opt1"],
                        params["erotype_opt2"],
                        params["init_moho_depth"],
                        params["final_moho_depth"],
                        params["ap_rad"],
                        params["ap_uranium"],
                        params["ap_thorium"],
                        params["zr_rad"],
                        params["zr_uranium"],
                        params["zr_thorium"],
                    )
                )
            failed += 1

    # Print warning(s) if more than one observed age of a given type was provided
    if (
        (len(params["obs_ahe"]) > 1)
        or (len(params["obs_aft"]) > 1)
        or (len(params["obs_zhe"]) > 1)
    ):
        print("")
        if len(params["obs_ahe"]) > 1:
            print(
                "WARNING: More than one measured AHe age supplied, only the first was written to the output file!"
            )
        if len(params["obs_aft"]) > 1:
            print(
                "WARNING: More than one measured AFT age supplied, only the first was written to the output file!"
            )
        if len(params["obs_zhe"]) > 1:
            print(
                "WARNING: More than one measured ZHe age supplied, only the first was written to the output file!"
            )

    print(
        "\n--- Execution complete ({0} succeeded, {1} failed) ---".format(
            success, failed
        )
    )

def run_model(params):
    # Say hello
    if not params["batch_mode"]:
        print("")
        print(30 * "-" + " Execution started " + 31 * "-")

    # Set flags if using batch mode
    if params["batch_mode"]:
        params["echo_info"] = False
        params["echo_thermal_info"] = False
        params["echo_ages"] = False
        params["plot_results"] = False

    # Conversion factors and unit conversions
    max_depth = kilo2base(params["max_depth"])
    moho_depth_init = kilo2base(params["init_moho_depth"])
    moho_depth = moho_depth_init
    delta_moho = kilo2base(params["init_moho_depth"] - params["final_moho_depth"])

    t_total = myr2sec(params["t_total"])
    dt = yr2sec(params["dt"])

    vx_init = mmyr2ms(params["vx_init"])
    crustal_flux = mmyr2ms(params["crustal_flux"])
    vx = calculate_erosion_rate(
        t_total,
        0.0,
        delta_moho,
        params["erotype"],
        params["erotype_opt1"],
        params["erotype_opt2"],
    )

    # Set number of passes needed based on erosion model type
    # Types 1-4 need only 1 pass
    if params["erotype"] < 5:
        num_pass = 1

    t_plots = myr2sec(np.array(params["t_plots"]))
    t_plots.sort()
    if len(t_plots) > 0:
        more_plots = True
    else:
        more_plots = False

    # Determine thickness of mantle to remove
    mantle_lith_thickness = max_depth - moho_depth
    removal_thickness = params["removal_fraction"] * mantle_lith_thickness

    # Calculate node spacing
    dx = max_depth / (params["nx"] - 1)  # m

    # Calculate time step
    nt = int(np.floor(t_total / dt))  # -

    # Echo model info if requested
    if params["echo_info"]:
        echo_model_info(
            dx,
            nt,
            dt,
            t_total,
            params["implicit"],
            vx,
            params["k_crust"],
            params["rho_crust"],
            params["cp_crust"],
            params["k_mantle"],
            params["rho_mantle"],
            params["cp_mantle"],
            params["k_a"],
            params["erotype"],
            cond_crit=0.5,
            adv_crit=0.5,
        )

    # Create arrays to hold temperature fields
    temp_new = np.zeros(params["nx"])
    temp_prev = np.zeros(params["nx"])

    # Create coordinates of the grid points
    x = np.linspace(0, max_depth, params["nx"])
    xstag = x[:-1] + dx / 2
    vx_hist = np.zeros(nt)
    depth_hist = np.zeros(nt)
    temp_hist = np.zeros(nt)
    time_hist = np.zeros(nt)
    if params["mantle_adiabat"]:
        adiabat_m = adiabat(
            alphav=params["alphav_mantle"],
            temp=params["temp_base"] + 273.15,
            cp=params["cp_mantle"],
        )
        temp_adiabat = params["temp_base"] + (xstag - max_depth) * adiabat_m
    else:
        adiabat_m = 0.0
        temp_adiabat = params["temp_base"]

    # Create material property arrays
    rho = np.ones(len(x)) * params["rho_crust"]
    rho[x > moho_depth] = params["rho_mantle"]
    cp = np.ones(len(x)) * params["cp_crust"]
    cp[x > moho_depth] = params["cp_mantle"]
    k = np.ones(len(xstag)) * params["k_crust"]
    k[xstag > moho_depth] = params["k_mantle"]
    heat_prod = np.ones(len(x)) * micro2base(params["heat_prod_crust"])
    heat_prod[x > moho_depth] = micro2base(params["heat_prod_mantle"])
    alphav = np.ones(len(x)) * params["alphav_crust"]
    alphav[x > moho_depth] = params["alphav_mantle"]

    # Generate initial temperature field
    if not params["batch_mode"]:
        print("")
        print("--- Calculating initial thermal model ---")
        print("")
    temp_init = temp_ss_implicit(
        params["nx"],
        dx,
        params["temp_surf"],
        params["temp_base"],
        vx_init,
        rho,
        cp,
        k,
        heat_prod,
    )
    interp_temp_init = interp1d(x, temp_init)
    init_moho_temp = interp_temp_init(moho_depth)
    init_heat_flow = kilo2base((k[0] + k[1]) / 2 * (temp_init[1] - temp_init[0]) / dx)
    if params["echo_thermal_info"]:
        print("- Initial surface heat flow: {0:.1f} mW/m^2".format(init_heat_flow))
        print("- Initial Moho temperature: {0:.1f}°C".format(init_moho_temp))
        print("- Initial Moho depth: {0:.1f} km".format(params["init_moho_depth"]))
        print(
            "- Initial LAB depth: {0:.1f} km".format(
                (max_depth - removal_thickness) / kilo2base(1)
            )
        )
        print("- Crustal flux: {0:.1f} mm/yr".format(crustal_flux / mmyr2ms(1)))

    # Create arrays to store elevation history
    elev_list = []
    time_list = []
    elev_list.append(0.0)
    time_list.append(0.0)

    # Calculate initial densities
    rho_prime = -rho * alphav * temp_init
    rho_inc_temp = rho + rho_prime

    # --- Set temperatures at 0 Ma ---
    temp_prev = temp_init.copy()

    # Modify temperatures and material properties for erotype 4
    if params["erotype"] == 4:
        temp_prev, moho_depth, rho, cp, k, heat_prod, alphav = init_erotype4(params, temp_init, x, xstag, temp_prev, moho_depth)

    # Modify temperatures for delamination
    if params["removal_fraction"] > 0.0:
        for ix in range(params["nx"]):
            if x[ix] > (max_depth - removal_thickness):
                temp_prev[ix] = params["temp_base"] + (x[ix] - max_depth) * adiabat_m

    if params["plot_results"]:
        # Set plot style
        plt.style.use("seaborn-darkgrid")

        # Plot initial temperature field
        fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 8))
        if t_plots.max() < t_total - 1.0:
            # Add an extra color for the final temperature if it is not in the
            # list of times for plotting
            colors = plt.cm.viridis_r(np.linspace(0, 1, len(t_plots) + 1))
        else:
            colors = plt.cm.viridis_r(np.linspace(0, 1, len(t_plots)))
        ax1.plot(temp_init, -x / 1000, "k:", label="Initial")
        ax1.plot(temp_prev, -x / 1000, "k-", label="0 Myr")
        ax2.plot(rho_inc_temp, -x / 1000, "k-", label="0 Myr")

    # Loop over number of required passes
    for j in range(num_pass):
        # Start the loop over time steps
        curtime = 0.0
        idx = 0
        if j == num_pass - 1:
            plotidx = 0

        # Restore values if using more than 1 pass
        if num_pass > 1:
            # Reset moho depth
            moho_depth = moho_depth_init
            # Reset initial temperatures
            for ix in range(params["nx"]):
                if x[ix] > (max_depth - removal_thickness):
                    temp_prev[ix] = params["temp_base"] + (x[ix] - max_depth) * adiabat_m
                else:
                    temp_prev[ix] = temp_init[ix]

        # Reset erosion rate
        vx = calculate_erosion_rate(
            t_total,
            curtime,
            delta_moho,
            params["erotype"],
            params["erotype_opt1"],
            params["erotype_opt2"],
        )

        # Calculate initial densities
        rho, cp, k, heat_prod, lab_depth = update_materials(
            x,
            xstag,
            moho_depth,
            params["rho_crust"],
            params["rho_mantle"],
            rho,
            params["cp_crust"],
            params["cp_mantle"],
            cp,
            params["k_crust"],
            params["k_mantle"],
            k,
            micro2base(params["heat_prod_crust"]),
            micro2base(params["heat_prod_mantle"]),
            heat_prod,
            temp_adiabat,
            temp_prev,
            params["k_a"],
            params["removal_fraction"],
        )
        rho_prime = -rho * alphav * temp_init
        rho_inc_temp = rho + rho_prime
        isoref = rho_inc_temp.sum() * dx
        h_ref = isoref / params["rho_a"]
        elev_init = max_depth - h_ref

        # Find starting depth if using only a one-pass erosion type
        if num_pass == 1:
            while curtime < t_total:
                curtime += dt
                vx_hist[idx] = vx
                idx += 1
                vx = calculate_erosion_rate(
                    t_total,
                    curtime,
                    delta_moho,
                    params["erotype"],
                    params["erotype_opt1"],
                    params["erotype_opt2"],
                )
            depth = (vx_hist * dt).sum()

            # Reset loop variables
            curtime = 0.0
            idx = 0
            vx = calculate_erosion_rate(
                t_total,
                curtime,
                delta_moho,
                params["erotype"],
                params["erotype_opt1"],
                params["erotype_opt2"],
            )

        if not params["batch_mode"]:
            print("")
            print(
                "--- Calculating transient thermal model (Pass {0}/{1}) ---".format(
                    j + 1, num_pass
                )
            )
            print("")
        while curtime < t_total:
            # if (idx+1) % 100 == 0:
            if not params["batch_mode"]:
                # print('- Step {0:5d} of {1} ({2:3d}%)\r'.format(idx+1, nt, int(round(100*(idx+1)/nt, 0))), end="")
                print(
                    "- Step {0:5d} of {1} (Time: {2:5.1f} Myr, Erosion rate: {3:5.2f} mm/yr)\r".format(
                        idx + 1, nt, curtime / myr2sec(1), vx / mmyr2ms(1)
                    ),
                    end="",
                )
            curtime += dt

            rho, cp, k, heat_prod, lab_depth = update_materials(
                x,
                xstag,
                moho_depth,
                params["rho_crust"],
                params["rho_mantle"],
                rho,
                params["cp_crust"],
                params["cp_mantle"],
                cp,
                params["k_crust"],
                params["k_mantle"],
                k,
                micro2base(params["heat_prod_crust"]),
                micro2base(params["heat_prod_mantle"]),
                heat_prod,
                temp_adiabat,
                temp_prev,
                params["k_a"],
                params["removal_fraction"],
            )
            if params["implicit"]:
                temp_new[:] = temp_transient_implicit(
                    params["nx"],
                    dx,
                    dt,
                    temp_prev,
                    params["temp_surf"],
                    params["temp_base"],
                    vx,
                    rho,
                    cp,
                    k,
                    heat_prod,
                )
            else:
                temp_new[:] = temp_transient_explicit(
                    temp_prev,
                    temp_new,
                    params["temp_surf"],
                    params["temp_base"],
                    params["nx"],
                    dx,
                    vx,
                    dt,
                    rho,
                    cp,
                    k,
                    heat_prod,
                )

            temp_prev[:] = temp_new[:]

            rho_prime = -rho * alphav * temp_new
            rho_temp_new = rho + rho_prime

            # Blend materials when the Moho lies between two nodes
            isonew = 0.0
            for i in range(len(rho_temp_new) - 1):
                rho_inc = rho_temp_new[i]
                if (moho_depth < x[i + 1]) and (moho_depth >= x[i]):
                    crust_frac = (moho_depth - x[i]) / dx
                    mantle_frac = 1.0 - crust_frac
                    rho_inc = (
                        crust_frac * rho_temp_new[i] + mantle_frac * rho_temp_new[i + 1]
                    )
                isonew += rho_inc * dx

            h_asthenosphere = isonew / params["rho_a"]
            elev = max_depth - h_asthenosphere

            # Update Moho depth
            moho_depth -= (vx - crustal_flux) * dt

            # Store tracked surface elevations and advection velocities
            if j == 0:
                elev_list.append(elev - elev_init)
                time_list.append(curtime / myr2sec(1.0))

            # Save Temperature-depth history
            if j == num_pass - 1:
                # Store temperature, time, depth
                interp_temp_new = interp1d(x, temp_new)
                depth -= vx * dt
                depth_hist[idx] = depth
                time_hist[idx] = curtime
                if abs(depth) <= 1e-6:
                    temp_hist[idx] = 0.0
                else:
                    temp_hist[idx] = interp_temp_new(depth)

            # Update index
            idx += 1

            # Update erosion rate
            vx = calculate_erosion_rate(
                t_total,
                curtime,
                delta_moho,
                params["erotype"],
                params["erotype_opt1"],
                params["erotype_opt2"],
            )

            if j == num_pass - 1:
                if params["plot_results"] and more_plots:
                    if curtime > t_plots[plotidx]:
                        ax1.plot(
                            temp_new,
                            -x / 1000,
                            "-",
                            label="{0:.1f} Myr".format(t_plots[plotidx] / myr2sec(1)),
                            color=colors[plotidx],
                        )
                        ax2.plot(
                            rho_temp_new,
                            -x / 1000,
                            label="{0:.1f} Myr".format(t_plots[plotidx] / myr2sec(1)),
                            color=colors[plotidx],
                        )
                        if plotidx == len(t_plots) - 1:
                            more_plots = False
                        plotidx += 1
                        # tplot = t_plots[plotidx]

        if not params["batch_mode"]:
            print("")

    rho_prime = -rho * alphav * temp_new
    rho_temp_new = rho + rho_prime
    isonew = rho_temp_new.sum() * dx

    interp_temp_new = interp1d(x, temp_new)
    final_moho_temp = interp_temp_new(moho_depth)
    final_heat_flow = kilo2base((k[0] + k[1]) / 2 * (temp_new[1] - temp_new[0]) / dx)

    if not params["batch_mode"]:
        print("")

    if params["echo_thermal_info"]:
        print("")
        print("--- Final thermal model values ---")
        print("")
        print("- Final surface heat flow: {0:.1f} mW/m^2".format(final_heat_flow))
        print("- Final Moho temperature: {0:.1f}°C".format(final_moho_temp))
        print(
            "- Final Moho depth: {0:.1f} km ({1:+.1f} km from crustal flux)".format(
                moho_depth / kilo2base(1),
                (crustal_flux / mmyr2ms(1)) * t_total / myr2sec(1),
            )
        )
        print("- Final LAB depth: {0:.1f} km".format(lab_depth / kilo2base(1)))

    if params["calc_ages"]:
        # INPUT
        # time_i:the time values (in Myr) in descending order at which the thermal history
        # is given (ex: 100,50,20,10,0); the last value should always be 0; the first value
        # should be smaller than 1000.
        # temp_i: the thermal history in degree Celsius
        # n: the number of time-temperature pairs used  to describe the temperature history
        # out_flag:  =0 only calculate fission track age
        #            =1 also calculate track length distribution and statistics
        # param_flag : =1 uses Laslett et al, 1987 parameters
        #              =2 uses Crowley et al., 1991 Durango parameters
        #              =3 uses Crowley et al., 1991 F-apatite parameters
        time_ma = t_total - time_hist
        time_ma = time_ma / myr2sec(1)

        if params["madtrax"]:
            age, _, _, _ = Mad_Trax(time_ma, temp_hist, len(time_ma), 1, 2)

        # Write time-temperature history to file for (U-Th)/He age prediction
        with open("time_temp_hist.csv", "w") as csvfile:
            writer = csv.writer(csvfile, delimiter=",", lineterminator="\n")
            # Write time-temperature history in reverse order!
            for i in range(-1, -(len(time_ma) + 1), -100):
                writer.writerow([time_ma[i], temp_hist[i]])

            # Write fake times if time history padding is enabled
            if params["pad_thist"]:
                if params["pad_time"] > 0.0:
                    # Make array of pad times with 1.0 Myr time increments
                    pad_times = np.arange(
                        t_total / myr2sec(1),
                        t_total / myr2sec(1) + params["pad_time"] + 0.1,
                        1.0,
                    )
                    for pad_time in pad_times:
                        writer.writerow([pad_time, temp_hist[i]])

        # Write time-temperature-depth history to file for reference
        with open("time_temp_depth_hist.csv", "w") as csvfile:
            writer = csv.writer(csvfile, delimiter=",", lineterminator="\n")
            # Write header
            writer.writerow(["Time (Ma)", "Temperature (C)", "Depth (m)"])
            # Write time-temperature history in reverse order!
            for i in range(-1, -(len(time_ma) + 1), -100):
                writer.writerow([time_ma[i], temp_hist[i], depth_hist[i]])

        ahe_age, corr_ahe_age, zhe_age, corr_zhe_age = he_ages(
            file="time_temp_hist.csv",
            ap_rad=params["ap_rad"],
            ap_uranium=params["ap_uranium"],
            ap_thorium=params["ap_thorium"],
            zr_rad=params["zr_rad"],
            zr_uranium=params["zr_uranium"],
            zr_thorium=params["zr_thorium"],
        )
        if params["ketch_aft"]:
            aft_age, aft_mean_ftl = ft_ages("time_temp_hist.csv")

        # Find effective closure temperatures
        ahe_temp = np.interp(float(corr_ahe_age), np.flip(time_ma), np.flip(temp_hist))
        aft_temp = np.interp(float(aft_age), np.flip(time_ma), np.flip(temp_hist))
        zhe_temp = np.interp(float(corr_zhe_age), np.flip(time_ma), np.flip(temp_hist))

        if params["batch_mode"]:
            tt_filename = params["model_id"] + "-time_temp_hist.csv"
            ttd_filename = params["model_id"] + "-time_temp_depth_hist.csv"
            ftl_filename = params["model_id"] + "-ft_length.csv"
            os.rename("time_temp_hist.csv", "batch_output/" + tt_filename)
            os.rename("time_temp_depth_hist.csv", "batch_output/" + ttd_filename)
            os.rename("ft_length.csv", "batch_output/" + ftl_filename)

        if params["echo_ages"]:
            print("")
            print("--- Predicted thermochronometer ages ---")
            print("")
            print(
                "- AHe age: {0:.2f} Ma (uncorrected age: {1:.2f} Ma)".format(
                    float(corr_ahe_age), float(ahe_age)
                )
            )
            if params["madtrax"]:
                print("- AFT age: {0:.2f} Ma (MadTrax)".format(age / 1e6))
            if params["ketch_aft"]:
                print("- AFT age: {0:.2f} Ma (Ketcham)".format(float(aft_age)))
            print(
                "- ZHe age: {0:.2f} Ma (uncorrected age: {1:.2f} Ma)".format(
                    float(corr_zhe_age), float(zhe_age)
                )
            )

        # If measured ages have been provided, calculate misfit
        if len(params["obs_ahe"]) + len(params["obs_aft"]) + len(params["obs_zhe"]) > 0:
            # Create single arrays of ages for misfit calculation
            pred_ages = []
            obs_ages = []
            obs_stdev = []
            for i in range(len(params["obs_ahe"])):
                pred_ages.append(float(corr_ahe_age))
                obs_ages.append(params["obs_ahe"][i])
                obs_stdev.append(params["obs_ahe_stdev"][i])
            for i in range(len(params["obs_aft"])):
                pred_ages.append(float(aft_age))
                obs_ages.append(params["obs_aft"][i])
                obs_stdev.append(params["obs_aft_stdev"][i])
            for i in range(len(params["obs_zhe"])):
                pred_ages.append(float(corr_zhe_age))
                obs_ages.append(params["obs_zhe"][i])
                obs_stdev.append(params["obs_zhe_stdev"][i])

            # Convert lists to NumPy arrays
            pred_ages = np.array(pred_ages)
            obs_ages = np.array(obs_ages)
            obs_stdev = np.array(obs_stdev)

            # Calculate misfit
            misfit = calculate_misfit(
                pred_ages,
                obs_ages,
                obs_stdev,
                params["misfit_num_params"],
                params["misfit_type"],
            )

            # Print misfit to the screen
            if params["echo_ages"]:
                print("")
                print("--- Predicted and observed age misfit ---")
                print("")
                print(
                    "- Misfit: {0:.4f} (misfit type {1}, {2} age(s))".format(
                        misfit, params["misfit_type"], len(pred_ages)
                    )
                )

    if (
        (params["plot_results"] and params["save_plots"])
        or params["write_temps"]
        or params["read_temps"]
    ):
        fp = "/Users/whipp/Work/Research/projects/Kellett-Coutand-Canadian-Cordillera/TC1D/"

    if params["plot_results"]:
        # Plot the final temperature field
        xmin = 0.0
        # xmax = params['temp_base'] + 100
        xmax = 1600.0
        ax1.plot(
            temp_new,
            -x / 1000,
            "-",
            label="{0:.1f} Myr".format(curtime / myr2sec(1)),
            color=colors[-1],
        )
        ax1.plot(
            [xmin, xmax],
            [-moho_depth / kilo2base(1), -moho_depth / kilo2base(1)],
            linestyle="--",
            color="black",
            lw=0.5,
        )
        ax1.plot(
            [xmin, xmax],
            [-params["init_moho_depth"], -params["init_moho_depth"]],
            linestyle="--",
            color="gray",
            lw=0.5,
        )

        if params["crust_solidus"]:
            crust_slice = x / 1000.0 <= params["final_moho_depth"]
            pressure = calculate_pressure(rho_temp_new, dx)
            crustal_pressure = pressure[crust_slice]
            crust_solidus = calculate_crust_solidus(
                params["crust_solidus_comp"], crustal_pressure
            )
            ax1.plot(
                crust_solidus,
                -x[crust_slice] / 1000.0,
                color="gray",
                linestyle=":",
                lw=1.5,
                label="Crust solidus",
            )

        if params["mantle_solidus"]:
            mantle_slice = x / 1000 >= params["final_moho_depth"]
            pressure = calculate_pressure(rho_temp_new, dx)
            mantle_solidus = calculate_mantle_solidus(
                pressure / 1.0e9, xoh=params["mantle_solidus_xoh"]
            )
            ax1.plot(
                mantle_solidus[mantle_slice],
                -x[mantle_slice] / 1000,
                color="gray",
                linestyle="--",
                lw=1.5,
                label="Mantle solidus",
            )

        ax1.text(20.0, -moho_depth / kilo2base(1) + 1.0, "Final Moho")
        ax1.text(20.0, -params["init_moho_depth"] - 3.0, "Initial Moho", color="gray")
        ax1.legend()
        ax1.axis([xmin, xmax, -max_depth / 1000, 0])
        ax1.set_xlabel("Temperature (°C)")
        ax1.set_ylabel("Depth (km)")
        # ax1.grid()

        xmin = 2700
        xmax = 3300
        ax2.plot(
            rho_temp_new,
            -x / 1000,
            label="{0:.1f} Myr".format(t_total / myr2sec(1)),
            color=colors[-1],
        )
        ax2.plot(
            [xmin, xmax],
            [-moho_depth / kilo2base(1), -moho_depth / kilo2base(1)],
            linestyle="--",
            color="black",
            lw=0.5,
        )
        ax2.plot(
            [xmin, xmax],
            [-params["init_moho_depth"], -params["init_moho_depth"]],
            linestyle="--",
            color="gray",
            lw=0.5,
        )
        ax2.axis([xmin, xmax, -max_depth / 1000, 0])
        ax2.set_xlabel("Density (kg m$^{-3}$)")
        ax2.set_ylabel("Depth (km)")
        ax2.legend()
        # ax2.grid()

        plt.tight_layout()
        if params["save_plots"]:
            plt.savefig(fp + "png/T_rho_hist.png", dpi=300)
        plt.show()

        fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(12, 8))
        # ax1.plot(time_list, elev_list, 'k-')
        ax1.plot(time_list, elev_list)
        ax1.set_xlabel("Time (Myr)")
        ax1.set_ylabel("Elevation (m)")
        ax1.set_xlim(0.0, t_total / myr2sec(1))
        ax1.set_title("Elevation history")
        # plt.axis([0.0, t_total/myr2sec(1), 0, 750])
        # ax1.grid()

        ax2.plot(time_hist / myr2sec(1), vx_hist / mmyr2ms(1))
        ax2.fill_between(
            time_hist / myr2sec(1),
            vx_hist / mmyr2ms(1),
            0.0,
            alpha=0.33,
            color="tab:blue",
            label="Erosion magnitude: {0:.1f} km".format(
                params["init_moho_depth"] - params["final_moho_depth"]
            ),
        )
        ax2.set_xlabel("Time (Myr)")
        ax2.set_ylabel("Erosion rate (mm/yr)")
        ax2.set_xlim(0.0, t_total / myr2sec(1))
        if params["erotype_opt1"] >= 0.0:
            ax2.set_ylim(ymin=0.0)
        # plt.axis([0.0, t_total/myr2sec(1), 0, 750])
        # ax2.grid()
        ax2.legend()

        plt.tight_layout()
        if params["save_plots"]:
            plt.savefig(fp + "png/elev_hist.png", dpi=300)
        plt.show()

        # fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(12, 8))

        # create objects
        fig = plt.figure(figsize=(12, 8))
        gs = GridSpec(3, 3, figure=fig)

        # create sub plots as grid
        ax1 = fig.add_subplot(gs[0:2, :])
        ax2 = fig.add_subplot(gs[2, :-1])
        ax3 = fig.add_subplot(gs[2, -1])
        # ax1.plot(time_ma, temp_hist, 'r-', lw=2)

        # Calculate synthetic uncertainties
        ahe_uncert = 0.1
        aft_uncert = 0.2
        zhe_uncert = 0.1
        ahe_min, ahe_max = (1.0 - ahe_uncert) * float(corr_ahe_age), (
            1.0 + ahe_uncert
        ) * float(corr_ahe_age)
        aft_min, aft_max = (1.0 - aft_uncert) * float(aft_age), (
            1.0 + aft_uncert
        ) * float(aft_age)
        zhe_min, zhe_max = (1.0 - zhe_uncert) * float(corr_zhe_age), (
            1.0 + zhe_uncert
        ) * float(corr_zhe_age)
        ax1.plot(time_ma, temp_hist)

        # Plot shaded uncertainty area and AHe age if no measured ages exist
        if len(params["obs_ahe"]) == 0:
            ax1.axvspan(
                ahe_min,
                ahe_max,
                alpha=0.33,
                color="tab:blue",
                label="Predicted AHe age ({0:.2f} Ma ± {1:.0f}% uncertainty; T$_c$ = {2:.1f}°C)".format(
                    float(corr_ahe_age), ahe_uncert * 100.0, ahe_temp
                ),
            )
            ax1.plot(float(corr_ahe_age), ahe_temp, marker="o", color="tab:blue")
        # Plot predicted age + observed AHe age(s)
        else:
            ax1.scatter(
                float(corr_ahe_age),
                ahe_temp,
                marker="o",
                color="tab:blue",
                label="Predicted AHe age ({0:.2f} Ma; T$_c$ = {1:.1f}°C)".format(
                    float(corr_ahe_age), ahe_temp
                ),
            )
            ahe_temps = []
            for i in range(len(params["obs_ahe"])):
                ahe_temps.append(ahe_temp)
            ax1.errorbar(
                params["obs_ahe"],
                ahe_temps,
                xerr=params["obs_ahe_stdev"],
                marker="s",
                color="tab:blue",
                label="Measured AHe age(s)",
            )

        # Plot shaded uncertainty area and AFT age if no measured ages exist
        if len(params["obs_aft"]) == 0:
            ax1.axvspan(
                aft_min,
                aft_max,
                alpha=0.33,
                color="tab:orange",
                label="Predicted AFT age ({0:.2f} Ma ± {1:.0f}% uncertainty; T$_c$ = {2:.1f}°C)".format(
                    float(aft_age), aft_uncert * 100.0, aft_temp
                ),
            )
            ax1.plot(float(aft_age), aft_temp, marker="o", color="tab:orange")
        # Plot predicted age + observed AFT age(s)
        else:
            ax1.scatter(
                float(aft_age),
                aft_temp,
                marker="o",
                color="tab:orange",
                label="Predicted AFT age ({0:.2f} Ma; T$_c$ = {1:.1f}°C)".format(
                    float(aft_age), aft_temp
                ),
            )
            aft_temps = []
            for i in range(len(params["obs_aft"])):
                aft_temps.append(aft_temp)
            ax1.errorbar(
                params["obs_aft"],
                aft_temps,
                xerr=params["obs_aft_stdev"],
                marker="s",
                color="tab:orange",
                label="Measured AFT age(s)",
            )

        # Plot shaded uncertainty area and ZHe age if no measured ages exist
        if len(params["obs_zhe"]) == 0:
            ax1.axvspan(
                zhe_min,
                zhe_max,
                alpha=0.33,
                color="tab:green",
                label="Predicted ZHe age ({0:.2f} Ma ± {1:.0f}% uncertainty; T$_c$ = {2:.1f}°C)".format(
                    float(corr_zhe_age), zhe_uncert * 100.0, zhe_temp
                ),
            )
            ax1.plot(float(corr_zhe_age), zhe_temp, marker="o", color="tab:green")
        # Plot predicted age + observed ZHe age(s)
        else:
            ax1.scatter(
                float(corr_zhe_age),
                zhe_temp,
                marker="o",
                color="tab:green",
                label="Predicted ZHe age ({0:.2f} Ma; T$_c$ = {1:.1f}°C)".format(
                    float(corr_zhe_age), zhe_temp
                ),
            )
            zhe_temps = []
            for i in range(len(params["obs_zhe"])):
                zhe_temps.append(zhe_temp)
            ax1.errorbar(
                params["obs_zhe"],
                zhe_temps,
                xerr=params["obs_zhe_stdev"],
                marker="s",
                color="tab:green",
                label="Measured ZHe age(s)",
            )

        ax1.set_xlim(t_total / myr2sec(1), 0.0)
        ax1.set_ylim(ymin=params["temp_surf"])
        ax1.set_xlabel("Time (Ma)")
        ax1.set_ylabel("Temperature (°C)")
        # Include misfit in title if there are measured ages
        if (
            len(params["obs_ahe"]) + len(params["obs_aft"]) + len(params["obs_zhe"])
            == 0
        ):
            ax1.set_title("Thermal history for surface sample")
        else:
            ax1.set_title(
                "Thermal history for surface sample (misfit = {0:.4f}; {1} age(s))".format(
                    misfit, len(obs_ages)
                )
            )
        if params["pad_thist"] and params["pad_time"] > 0.0:
            ax1.annotate(
                "Initial holding time: +{0:.1f} Myr".format(params["pad_time"]),
                xy=(time_ma.max(), temp_hist[0]),
                xycoords="data",
                xytext=(0.95 * time_ma.max(), 0.65 * temp_hist.max()),
                textcoords="data",
                arrowprops=dict(arrowstyle="->", connectionstyle="arc3", fc="black"),
                bbox=dict(boxstyle="round4,pad=0.3", fc="white", lw=0),
            )
        # ax1.grid()
        ax1.legend()

        ax2.plot(time_ma, vx_hist / mmyr2ms(1))
        ax2.fill_between(
            time_ma,
            vx_hist / mmyr2ms(1),
            0.0,
            alpha=0.33,
            color="tab:blue",
            label="Erosion magnitude: {0:.1f} km".format(
                params["init_moho_depth"] - params["final_moho_depth"]
            ),
        )
        ax2.set_xlabel("Time (Ma)")
        ax2.set_ylabel("Erosion rate (mm/yr)")
        ax2.set_xlim(t_total / myr2sec(1), 0.0)
        if params["erotype_opt1"] >= 0.0:
            ax2.set_ylim(ymin=0.0)
        # plt.axis([0.0, t_total/myr2sec(1), 0, 750])
        # ax2.grid()
        ax2.legend()
        ax2.set_title("Erosion history for surface sample")

        ft_lengths = np.genfromtxt("ft_length.csv", delimiter=",", skip_header=1)
        length = ft_lengths[:, 0]
        prob = ft_lengths[:, 1]
        ax3.plot(length, prob)
        ax3.plot(
            [float(aft_mean_ftl), float(aft_mean_ftl)],
            [0.0, 1.05 * prob.max()],
            label="Mean: {0:.1f} µm".format(float(aft_mean_ftl)),
        )
        ax3.set_xlabel("Track length (um)")
        ax3.set_ylabel("Probability")
        ax3.set_xlim([0.0, 20.0])
        ax3.set_ylim([0.0, 1.05 * prob.max()])
        ax3.legend()
        ax3.set_title("Apatite fission-track length distribution")

        plt.tight_layout()
        if params["save_plots"]:
            plt.savefig(fp + "png/cooling_hist.png", dpi=300)
        plt.show()

    if params["read_temps"]:
        load_file = "py/output_temps.csv"
        data = np.genfromtxt(fp + load_file, delimiter=",", skip_header=1)
        temps = data[:, 1]
        temp_diff = temps[1:] - temp_new[1:]
        pct_diff = temp_diff / temps[1:] * 100.0
        plt.figure(figsize=(12, 6))
        plt.plot(pct_diff, -x[1:] / 1000, "k-")
        plt.xlabel("Percent temperature difference")
        plt.ylabel("Depth (km)")
        plt.grid()
        plt.title("Percent difference from explicit FD solution")
        plt.show()

    if params["write_temps"]:
        print("")
        print("--- Writing temperature output to file ---")
        print("")
        temp_x_out = np.zeros([len(x), 2])
        temp_x_out[:, 0] = x
        temp_x_out[:, 1] = temp_new
        savefile = "py/output_temps.csv"
        np.savetxt(
            fp + savefile,
            temp_x_out,
            delimiter=",",
            header="Depth (m),Temperature(deg. C)",
        )
        print("- Temperature output saved to file\n  " + fp + savefile)

    if params["batch_mode"]:
        # Write output to a file
        outfile = "TC1D_batch_log.csv"

        # Define measured ages for batch output
        if len(params["obs_ahe"]) == 0:
            obs_ahe = -9999.0
            obs_ahe_stdev = -9999.0
        else:
            obs_ahe = params["obs_ahe"][0]
            obs_ahe_stdev = params["obs_ahe_stdev"][0]
        if len(params["obs_aft"]) == 0:
            obs_aft = -9999.0
            obs_aft_stdev = -9999.0
        else:
            obs_aft = params["obs_aft"][0]
            obs_aft_stdev = params["obs_aft_stdev"][0]
        if len(params["obs_zhe"]) == 0:
            obs_zhe = -9999.0
            obs_zhe_stdev = -9999.0
        else:
            obs_zhe = params["obs_zhe"][0]
            obs_zhe_stdev = params["obs_zhe_stdev"][0]

        # Define misfit details for output
        if (
            len(params["obs_ahe"]) + len(params["obs_aft"]) + len(params["obs_zhe"])
            == 0
        ):
            misfit = -9999.0
            misfit_type = -9999.0
            misfit_ages = 0
        else:
            misfit_type = params["misfit_type"]
            misfit_ages = len(obs_ages)

        # Open file for writing
        with open(outfile, "a+") as f:
            f.write(
                "{0:.4f},{1:.4f},{2:.4f},{3},{4:.4f},{5:.4},{6},{7:.4f},{8:.4f},"
                "{9},{10:.4f},{11:.4f},{12:.4f},{13:.4f},{14:.4f},{15:.4f},{16:.4f},"
                "{17:.4f},{18:.4f},{19:.4f},{20:.4f},{21:.4f},{22:.4f},{23:.4f},{24:.4f},"
                "{25:.4f},{26:.4f},{27:.4f},{28:.4f},{29:.4f},{30:.4f},{31:.4f},{32:.4f},{33:.4f},"
                "{34:.4f},{35:.4f},{36:.4f},{37:.4f},{38:.6f},{39},{40}"
                "\n".format(
                    t_total / myr2sec(1),
                    dt / yr2sec(1),
                    max_depth / kilo2base(1),
                    params["nx"],
                    params["temp_surf"],
                    params["temp_base"],
                    params["mantle_adiabat"],
                    params["rho_crust"],
                    params["removal_fraction"],
                    params["erotype"],
                    params["erotype_opt1"],
                    params["erotype_opt2"],
                    params["init_moho_depth"],
                    init_moho_temp,
                    init_heat_flow,
                    elev_list[1] / kilo2base(1),
                    params["final_moho_depth"],
                    final_moho_temp,
                    final_heat_flow,
                    elev_list[-1] / kilo2base(1),
                    params["ap_rad"],
                    params["ap_uranium"],
                    params["ap_thorium"],
                    params["zr_rad"],
                    params["zr_uranium"],
                    params["zr_thorium"],
                    float(corr_ahe_age),
                    ahe_temp,
                    obs_ahe,
                    obs_ahe_stdev,
                    float(aft_age),
                    aft_temp,
                    obs_aft,
                    obs_aft_stdev,
                    float(corr_zhe_age),
                    zhe_temp,
                    obs_zhe,
                    obs_zhe_stdev,
                    misfit,
                    misfit_type,
                    misfit_ages,
                )
            )

    if not params["batch_mode"]:
        print("")
        print(30 * "-" + " Execution complete " + 30 * "-")

max_stars_repo_name: HUGG/Dal-thermochron

max_issues_repo_name: HUGG/Dal-thermochron

max_forks_repo_name: HUGG/Dal-thermochron


text:
# -*- coding: utf-8 -*-
"""This module defines classes and functions for browsing QuartataWeb.

----------------------------------------------------------------------------------------
Based on code written by the CHARMM-GUI team (http://charmm-gui.org) and modified by James Krieger

This suite uses the following softwares:
a) python Splinter package (https://splinter.readthedocs.org/en/latest/)
b) a web browser, such as Google Chrome or Mozilla Firefox
c) the corresponding driver such as chromedriver (https://sites.google.com/a/chromium.org/chromedriver/downloads)
   for Chrome or geckodriver (https://github.com/mozilla/geckodriver/releases) for Firefox
----------------------------------------------------------------------------------------
"""

from prody import PY3K, LOGGER
import numpy as np

__all__ = ['QuartataWebBrowser', 'QuartataChemicalRecord', 'searchQuartataWeb']

class QuartataWebBrowser(object):
    """Class to browse the QuartataWeb website."""

    def __init__(self, data_source=None, drug_group=None, input_type=None, query_type=None,
                 data=None, num_predictions=None, browser_type=None, job_id=None):
        """Instantiate a QuartataWebBrowser object instance.

        :arg data_source: source database for QuartataWeb analysis
            options are ``"DrugBank"`` or ``"STITCH"``. Default is ``"DrugBank"``
        :type data_source: str

        :arg drug_group: group of drugs if using DrugBank
            options are ``"Approved"`` or ``"All"``. Default is ``"All"``
        :type drug_group: str

        :arg input_type: number corresponding to the input type, options are
            ``1`` (Chemical and/or target) or
            ``2`` (A list of chemicals, targets or chemical combinations).
            Default is ``1``
        :type input_type: int

        :arg query_type: number corresponding to the query type. Options are
            dependent on input_type.

            With input_type 1, they are:
            * ``1`` (chemical-target interaction)
            * ``2`` (chemical-chemical similarity)
            * ``3`` (target-target similarity)

            With input_type 2, they are:
            * ``1`` (chemicals)
            * ``2`` (targets)
            * ``3`` (chemical combinations)

            Default is ``1``
        :type query_type: int

        :arg data: data to enter into the box or boxes. This varies depending on input type
            and query type, but will always be a list of strings.

            For input_type 1, a list with two items is expected. These will be one of the
            following depending on query_type:
            * With query_type 1, the first would be a chemical and the second a target.
                One of these can also be left blank.
            * With query_type 2, the first would be a chemical and the second a chemical.
            * With query_type 3, the first would be a target and the second a target.

            For input_type 2, a list with any length is expected. These will be one of the
            following depending on query_type:
            * With query_type 1, these would be chemicals.
            * With query_type 2, these would be targets.
            * With query_type 3, these would be pairs of chemicals, separated by semicolons.
        :type data: list

        :arg num_predictions: number of predictions to show or consider in addition to
            known interactions. Default is ``0``.
            With DrugBank and input_type 1, a second number can be provided in a list
            for secondary interactions.
        :type num_predictions: int, list

        :arg browser_type: browser type for navigation
            Default is ``"Chrome"``
        :type browser_type: str

        :arg job_id: job ID for accessing previous jobs
            Default is ``None``
        :type browser_type: int
        """

        self.browser_type = None
        self.browser = None

        self.data_source = None
        self.drug_group = None
        self.input_type = None
        self.query_type = None
        self.data = None
        self.num_predictions = None

        self.job_id = job_id

        self.setBrowserType(browser_type)
        self.setDataSource(data_source)
        self.setDrugGroup(drug_group)
        self.setInputType(input_type)
        self.setQueryType(query_type)
        self.setData(data)
        self.setNumPredictions(num_predictions)

    def updateHomePage(self):
        """Update the home page with data from setting variables"""
        url = "http://quartata.csb.pitt.edu"

        if self.data_source == 'DrugBank':
            url += '/index'
        else:
            url += '/index_stitch'

        if self.input_type == 2:
            url += '_2'

        url += '.php'
        self.browser.visit(url)

        if self.data_source == 'DrugBank':
            if self.drug_group == 'Approved':
                self.browser.find_by_name('db_type')[0].click()
            else:
                self.browser.find_by_name('db_type')[1].click()

        if self.query_type is not None:
            self.browser.find_by_name('pattern')[self.query_type - 1].click()

        if self.data is not None:
            if self.input_type == 1:
                if self.query_type == 1:
                    self.browser.find_by_name('q_drug_1')[0].fill(self.data[0])
                    self.browser.find_by_name('q_target_1')[
                        0].fill(self.data[1])
                elif self.query_type == 2:
                    self.browser.find_by_name('q_drug_1')[0].fill(self.data[0])
                    self.browser.find_by_name('q_drug_2')[0].fill(self.data[1])
                else:
                    self.browser.find_by_name('q_target_1')[0].fill(self.data[0])
                    self.browser.find_by_name('q_target_2')[0].fill(self.data[1])
            else:
                if self.query_type == 1:
                    self.browser.find_by_name('q_drugs')[0].fill('\n'.join(self.data))
                if self.query_type == 2:
                    self.browser.find_by_name('q_targets')[0].fill('\n'.join(self.data))
                else:
                    self.browser.find_by_name('q_drug_pairs')[0].fill('\n'.join(self.data))

        if self.num_predictions is not None:
            self.browser.find_by_name('pred_n')[0].fill(self.num_predictions[0])
            if self.data_source == 'DrugBank' and self.input_type == 1:
                self.browser.find_by_name('pred_n_2nd')[0].fill(self.num_predictions[1])

    def setDataSource(self, data_source):
        """Set data_source and update home page

        :arg data_source: source database for QuartataWeb analysis
            options are ``"DrugBank"`` or ``"STITCH"``. Default is ``"DrugBank"``
        :type data_source: str
        """
        if data_source is None:
            data_source = 'DrugBank'
        elif not isinstance(data_source, str):
            raise TypeError('data_source should be a string or None')
        elif data_source.lower() == 'drugbank':
            data_source = 'DrugBank'
        elif data_source.lower() == 'stitch':
            data_source = 'STITCH'
        else:
            raise ValueError('data_source should be DrugBank, STITCH or None')

        self.data_source = data_source
        self.updateHomePage()

    def setDrugGroup(self, group):
        """Set drug_group and update home page

        :arg group: group of drugs if using DrugBank
            options are ``"Approved"`` or ``"All"``. Default is ``"All"``
        :type group: str
        """
        if self.data_source == 'DrugBank':
            if group is None:
                group = 'All'
            elif not isinstance(group, str):
                raise TypeError('group must be string or None')
            elif group.lower() == 'all':
                group = 'All'
            elif group.lower() == 'approved':
                group = 'Approved'
            else:
                raise ValueError('group should be approved, all or None')

            self.drug_group = group
            self.updateHomePage()
        elif group is not None:
            LOGGER.warn('there are no groups when using STITCH')

    def setInputType(self, input_type):
        """Set input_type and update home page

        :arg input_type: number corresponding to the input type, options are
            ``1`` (Chemical and/or target) or
            ``2`` (A list of chemicals, targets or chemical combinations).
            Default is ``1``
        :type input_type: int
        """
        if input_type is None:
            input_type = 1
        elif not isinstance(input_type, int):
            raise TypeError('input_type should be an integer (1 or 2) or None')
        elif not input_type in [1, 2]:
            raise ValueError('input_type should be 1, 2 or None')

        self.input_type = input_type
        self.updateHomePage()

    def setQueryType(self, query_type):
        """Set query_type and update home page

        :arg query_type: number corresponding to the query type. Options are
            dependent on input_type.

            With input_type 1, they are:
            * ``1`` (chemical-target interaction)
            * ``2`` (chemical-chemical similarity)
            * ``3`` (target-target similarity)

            With input_type 2, they are:
            * ``1`` (chemicals)
            * ``2`` (targets)
            * ``3`` (chemical combinations)

            Default is ``1``
        :type query_type: int
        """
        if query_type is None:
            query_type = 1
        elif not isinstance(query_type, int):
            raise TypeError(
                'query_type should be an integer (1, 2 or 3) or None')
        elif not query_type in [1, 2, 3]:
            raise ValueError('query_type should be 1, 2, 3 or None')

        self.query_type = query_type
        self.updateHomePage()

    def setData(self, data):
        """Set data and update home page

        :arg data: data to enter into the box or boxes. This varies depending on input type
            and query type, but will always be a list of strings.

            For input_type 1, a list with two items is expected. These will be one of the
            following depending on query_type:
            * With query_type 1, the first would be a chemical and the second a target.
                One of these can also be left blank.
            * With query_type 2, the first would be a chemical and the second a chemical.
            * With query_type 3, the first would be a target and the second a target.

            For input_type 2, a list with any length is expected. These will be one of the
            following depending on query_type:
            * With query_type 1, these would be chemicals.
            * With query_type 2, these would be targets.
            * With query_type 3, these would be pairs of chemicals, separated by semicolons.
        :type data: list
        """
        if data is None:
            LOGGER.warn('data is not set')
        elif not isinstance(data, list):
            raise TypeError('data should be a list')
        else:
            for item in data:
                if not isinstance(item, str):
                    raise TypeError('data should be a list of strings')

            if self.input_type == 1:
                if len(data) > 2:
                    raise ValueError(
                        'data can only have two values with input_type 1')

                if len(data) == 1:
                    data.append('')

            if self.input_type == 3:
                for item in data:
                    if item.find(';') == -1:
                        raise ValueError(
                            'each item in data must be a pair with ; as delimiter')

        self.data = data
        self.updateHomePage()

    def setNumPredictions(self, num_predictions):
        """Set num_predictions and update home page

        :arg num_predictions: number of predictions to show or consider in addition to
            known interactions. Default is ``0``.
            With DrugBank and input_type 1, a second number can be provided in a list
            for secondary interactions.
        :type num_predictions: int, list
        """
        if num_predictions is None:
            num_predictions = 0

        if not isinstance(num_predictions, (int, list)):
            raise TypeError(
                'num_predictions should be an integer, a list or None')

        if isinstance(num_predictions, int):
            num_predictions = [num_predictions, 0]

        if num_predictions[0] > 100:
            raise ValueError('1st num_predictions must be <= 100')

        if num_predictions[1] > 20:
            raise ValueError('2nd num_predictions must be <= 20')

        self.num_predictions = num_predictions
        self.updateHomePage()

    def setBrowserType(self, browser_type):
        """Set browser_type and update home page

        :arg browser_type: browser type for navigation
            Default is ``"Chrome"``
        :type browser_type: str
        """
        if self.no_data:
            try:
                from splinter import Browser
            except ImportError:
                raise ImportError('Browser module could not be imported. '
                                'install splinter package to solve the problem.')
            else:
                from selenium.webdriver.common.service import WebDriverException

            if browser_type is None:
                try:
                    browser = Browser('chrome')
                    url = "http://quartata.csb.pitt.edu"
                    browser.visit(url)
                except WebDriverException:
                    try:
                        browser = Browser('firefox')
                        url = "http://quartata.csb.pitt.edu"
                        browser.visit(url)
                    except WebDriverException:
                        raise ValueError(
                            'No web driver found for Chrome or Firefox. Please specify a browser type or download an appropriate driver.')
                    else:
                        self.browser_type = 'firefox'
                else:
                    self.browser_type = 'chrome'

        elif not isinstance(browser_type, str):
            raise TypeError('browser_type should be a string or None')
        else:
            try:
                browser = Browser(browser_type)
                url = "http://quartata.csb.pitt.edu"
                browser.visit(url)
            except WebDriverException:
                raise ValueError(
                    'No web driver found for browser_type. Please specify a different browser type or download an appropriate driver.')
            else:
                self.browser_type = browser_type

        self.browser = browser
        self.updateHomePage()

    def setJObID(self, job_id):
        """Set job_id and view results

        :arg job_id: job ID for accessing previous jobs
            Default is ``None``
        :type browser_type: int
        """
        self.job_id = job_id
        self.viewResults()

    def viewResults(self):
        """View results by clicking submit or using a job_id"""
        if self.job_id is None or self.browser.url.find('index') != -1:
            self.browser.find_by_name('submit')[0].click()
            self.job_id = self.browser.url.split('_')[-1].split('=')[-1]

        else:
            if self.data_source == 'DrugBank':
                url = '_'.join(['http://quartata.csb.pitt.edu/quartata_result.php?job',
                                'id={0}'.format(self.job_id)])
            else:
                url = '_'.join(['http://quartata.csb.pitt.edu/quartata_result',
                                'stitch.php?job', 'id={0}'.format(self.job_id)])

            self.browser.visit(url)

    def goToDownloads(self):
        """Go to downloads page"""
        if self.job_id is None:
            self.viewResults()

        if self.data_source == 'DrugBank':
            url = '_'.join(['http://quartata.csb.pitt.edu/quartata_download.php?job',
                            'id={0}'.format(self.job_id)])
        else:
            url = '_'.join(['http://quartata.csb.pitt.edu/quartata_download',
                            'stitch.php?job', 'id={0}'.format(self.job_id)])

        self.browser.visit(url)

    def goToWorkDir(self):
        """Go to working directory"""
        if self.job_id is None:
            self.viewResults()

        url = 'http://quartata.csb.pitt.edu/work/{0}'.format(self.job_id)
        self.browser.visit(url)

    def parseChemicals(self):
        """Go to working directory and parse chemicals for query protein.
        Updates self.chemical_data"""
        try:
            self.goToWorkDir()

            if self.data_source == 'DrugBank':
                filename = 'known_drugs_for_query_protein.txt'
            else:
                filename = 'known_chemicals_for_query_protein.txt'

            self.browser.find_by_text(filename)[0].click()

            import requests
            html = requests.get(self.browser.url).content
            if PY3K:
                html = html.decode()

            lines = html.split('\n')

            self.fields = lines[0].split('\t')
            self.num_fields = len(self.fields)

            self.num_rows = len(lines[1:])
            if lines[-1].strip() == '':
                self.num_rows -= 1

            dtypes = []
            for i, item in enumerate(lines[1].split('\t')):
                if item.isnumeric():
                    dtypes.append((self.fields[i], int))
                elif item.find('.') != -1 and item.replace('.','0').isnumeric():
                    dtypes.append((self.fields[i], float))
                else:
                    dtypes.append((self.fields[i], object))

            self.chemical_data = np.empty(self.num_rows, dtype=dtypes)

            for i, line in enumerate(lines[1:self.num_rows+1]):
                items = line.split('\t')
                if len(items) != self.num_fields:
                    raise ValueError('line {0} has the wrong number of fields'.format(i+1))

                for j, item in enumerate(items):
                    self.chemical_data[i][j] = item
        except:
            success = False
        else:
            success = True
        return success

    def quit(self):
        self.browser.quit()

class QuartataChemicalRecord(object):
    """Class for handling chemical data from QuartataWebBrowser"""

    def __init__(self, data_source=None, drug_group=None, input_type=None, query_type=None,
                 data=None, num_predictions=None, browser_type=None, job_id=None):
        """Instantiate a QuartataChemicalRecord object instance.
        Inputs are the same as QuartataWebBrowser.
        """
        self._chemData = None
        self._filterDict = None
        self.data_source = data_source
        self.drug_group = drug_group
        self.input_type = input_type
        self.query_type = query_type
        self.data = data
        self.num_predictions = num_predictions
        self.browser_type = browser_type
        self.job_id = job_id

        self.isSuccess = self.fetch(data_source, drug_group, input_type, query_type,
                                    data, num_predictions, browser_type, job_id)

    def fetch(self, data_source=None, drug_group=None, input_type=None, query_type=None,
              data=None, num_predictions=None, browser_type=None, job_id=None):
        """Fetch data"""
        if data_source is None:
            data_source = self.data_source
        if drug_group is None:
            drug_group = self.drug_group
        if input_type is None:
            input_type = self.input_type
        if query_type is None:
            query_type = self.query_type
        if data is None:
            data = self.data
        if num_predictions is None:
            num_predictions = self.num_predictions
        if browser_type is None:
            browser_type = self.browser_type
        if job_id is None:
            job_id = self.job_id

        self.qwb = QuartataWebBrowser(data_source, drug_group, input_type, query_type,
                                      data, num_predictions, browser_type, job_id)

        isSuccess = self.qwb.parseChemicals()
        self.qwb.quit()

        self._chemData = self.qwb.chemical_data
        chem_temp_dict = dict()
        listAll = []
        for temp in self._chemData:
            temp_dict = dict()
            chem_name = temp[1]

            temp_dict['DB_ID'] = temp[0]
            temp_dict['chemical_name'] = chem_name
            temp_dict['mol_weight'] = temp[2]
            temp_dict['SMILES'] = temp[3]
            temp_dict['conf_score'] = temp[4]

            chem_temp_dict[chem_name] = temp_dict
            listAll.append(chem_name)

        self._listAll = tuple(listAll)
        self._list = self._listAll
        self._chemDict = chem_temp_dict

        return isSuccess

    def getChemicalList(self, filtered=True):
        """Returns chemical list (filters may be applied)"""
        if not self.isSuccess:
            LOGGER.warn('Quartata Chemical Record does not have any data yet.'
                        'Please run fetch again, possibly with different parameters.')

        if filtered:
            return self._list
        return self._listAll

    def getSMILESList(self, filtered=True):
        """Returns SMILES list (filters may be applied)"""
        if not self.isSuccess:
            LOGGER.warn('Quartata Chemical Record does not have any data yet.'
                        'Please run fetch again, possibly with different parameters.')

        if filtered:
            return [self._chemDict[key]['SMILES'] for key in self._list]
        return self._chemData['SMILES']

    def getParticularSMILES(self, key):
        """Returns SMILES for a particular chemical"""
        if not self.isSuccess:
            LOGGER.warn('Quartata Chemical Record does not have any data yet.'
                        'Please run fetch again, possibly with different parameters.')

        return self._chemDict[key]['SMILES']

    def getFilterList(self):
        """Returns a list of chemicals for the entries that were filtered out"""

        filterDict = self._filterDict
        if filterDict is None:
            raise ValueError('You cannot obtain the list of filtered out entries before doing any filtering.')

        temp_str = ', '.join([str(len(filterDict['lower_MW'])), str(len(filterDict['upper_MW'])),
                              str(len(filterDict['conf_score']))])
        LOGGER.info('Filtered out [' + temp_str + '] for [lower weight, upper weight, confidence score]')
        return self._filterList

    def filter(self, lower_weight=None, upper_weight=None, cutoff_score=None):
        """Filters out chemicals from the list and returns the updated list.
        Chemicals that satisfy any of the following criterion will be filtered out.
        (1) Molecular weight < lower_weight (must be a positive number);
        (2) Molecular weight > upper_weight (must be a positive number);
        (3) Confidence score < cutoff_score (must be a positive number);

        Please note that every time this function is run, this overrides any previous runs.
        Therefore, please provide all filters at once.
        """
        if not self.isSuccess:
            LOGGER.warn('Quartata Chemical Record does not have any data yet.'
                        'Please run fetch again, possibly with different parameters.')
            return None

        if lower_weight == None:
            lower_weight = 0
        elif not isinstance(lower_weight, (float, int)):
            raise TypeError('lower_weight must be a float or an integer')
        if lower_weight >= 0:
            lower_weight = float(lower_weight)
        else:
            raise ValueError('lower_weight must be a number not less than 0')

        if upper_weight == None:
            upper_weight = 0
        elif not isinstance(upper_weight, (float, int)):
            raise TypeError('upper_weight must be a float or an integer')
        if upper_weight >= 0:
            upper_weight = float(upper_weight)
        else:
            raise ValueError('upper_weight must be a number not less than 0')

        if cutoff_score == None:
            cutoff_score = 0
        elif not isinstance(cutoff_score, (float, int)):
            raise TypeError('cutoff_score must be a float or an integer')
        elif cutoff_score >= 0:
            cutoff_score = float(cutoff_score)
        else:
            raise ValueError('cutoff_score must be a number not less than 0')

        quartataInfo = self._chemDict
        if quartataInfo is None:
            raise ValueError("Quartata Chemical Record does not have any data yet. Please run fetch.")

        listAll = self._listAll
        ref_indices_set = set(range(self.qwb.num_rows))
        filterListLowerMW = []
        filterListUpperMW = []
        filterListConf = []

        for chem in listAll:
            temp_dict = quartataInfo[chem]

            if temp_dict['mol_weight'] < lower_weight:
                filterListLowerMW.append(chem)
                continue

            if upper_weight > 0 and temp_dict['mol_weight'] > upper_weight:
                filterListUpperMW.append(chem)
                continue

            if temp_dict['conf_score'] < cutoff_score:
                filterListConf.append(chem)
                continue

        filterList = filterListLowerMW + filterListUpperMW + filterListConf
        filterDict = {'lower_MW': filterListLowerMW, 'upper_MW': filterListUpperMW, 'conf_score': filterListConf}
        self._filterList = filterList
        self._filterDict = filterDict
        self._list = list(set(self._listAll) - set(filterList))
        LOGGER.info(str(len(self._listAll)-len(self._list)) + ' chemicals have been filtered out from '+str(len(self._listAll))+' QuartataWeb hits (remaining: '+str(len(self._list))+').')
        return self._list

def searchQuartataWeb(data_source=None, drug_group=None, input_type=None, query_type=None,
                   data=None, num_predictions=None, browser_type=None, job_id=None, result_type='Chemical'):
    """Wrapper function for searching QuartataWeb"""
    if result_type == 'Chemical':
        return QuartataChemicalRecord(data_source, drug_group, input_type, query_type,
                                    data, num_predictions, browser_type, job_id)
    else:
        LOGGER.warn('No other result types are supported yet')
        return None

searchQuartataWeb.__doc__ += QuartataWebBrowser.__doc__

ext: py

max_stars_repo_name: bwingert/ProDy

max_issues_repo_name: bwingert/ProDy

max_forks_repo_name: bwingert/ProDy

reason: import numpy


text:
#Karan Vombatkere
#German Enigma Machine
#October 2017

from string import *
import numpy as np

Letters = "ABCDEFGHIJKLMNOPQRSTUVWXYZ"

#function to create a dictionary with letters and their indices
#Use this to return the index of a letter (a = 0, .., z = 25)
def genDictionary():
    letter_index_pairs = []
    for indx, char in enumerate(Letters):
        letter_index_pairs.append([char, indx])

    Indx_Dict = dict(letter_index_pairs)
    print("Generated Letter Dictionary!")
    return Indx_Dict

#Call the function to create a global dictionary
Char_Indices = genDictionary()

#Enigma Rotor Configurations
Rotor1 = ['E','K','M','F','L','G','D','Q','V','Z','N','T','O','W','Y','H','X','U','S','P','A','I','B','R','C','J']
Rotor2 = ['A','J','D','K','S','I','R','U','X','B','L','H','W','T','M','C','Q','G','Z','N','P','Y','F','V','O','E']
Rotor3 = ['B','D','F','H','J','L','C','P','R','T','X','V','Z','N','Y','E','I','W','G','A','K','M','U','S','Q','O']
Rotor4 = ['E','S','O','V','P','Z','J','A','Y','Q','U','I','R','H','X','L','N','F','T','G','K','D','C','M','W','B']
Rotor5 = ['V','Z','B','R','G','I','T','Y','U','P','S','D','N','H','L','X','A','W','M','J','Q','O','F','E','C','K']
Rotor6 = ['J','P','G','V','O','U','M','F','Y','Q','B','E','N','H','Z','R','D','K','A','S','X','L','I','C','T','W']
Rotor7 = ['N','Z','J','H','G','R','C','X','M','Y','S','W','B','O','U','F','A','I','V','L','P','E','K','Q','D','T']
Rotor8 = ['F','K','Q','H','T','L','X','O','C','B','J','S','P','D','Z','R','A','M','E','W','N','I','U','Y','G','V']

#Set of All Rotors
Enigma_Rotors = [Rotor1, Rotor2, Rotor3, Rotor4, Rotor5, Rotor6, Rotor7, Rotor8]

#Class to implement Rotors
#Rotor is Initialized with a list of letters as given and can rotate the pointer to keep track of the position
class Rotor:
    'Implement a Rotor for the Enigma Machine'
    rotorPos = 0
    def __init__(self, rotorPermutation, pos):
        self.rotorPermutation = list(rotorPermutation)
        self.rotatebyN(pos)
        self.rotorPos = pos
        self.origPos = 0
        self.rotateFlag = False #All rotors are configured to NOT rotate initially
        #print("Rotor Configuration:\n", self.rotorPermutation)

    #Function to rotate the rotor by a specified initial rotation, n
    def rotatebyN(self, n):
        for i in range(n):
            self.rotate()

    #Function to rotate the rotor once
    def rotate(self):
        self.rotorPos += 1 #keep track of positions
        self.rotorPos = self.rotorPos % 26

        #shift all elements in rotorPermutation one space to the right
        rotatedRotor = []
        for i in range(len(self.rotorPermutation)):
            rotatedRotor.append(self.rotorPermutation[(i-1)%26])

        #increment each character by 1
        self.rotorPermutation = [Letters[((Char_Indices[ch]) +1)%26] for indx, ch in enumerate(list(rotatedRotor))]
        #print(self.rotorPermutation)

    #display the Details of the rotor and its current location
    def displayRotor(self):
        print("Rotor Configuration:", self.rotorPermutation)
        print("Rotor Position:", self.rotorPos, "\nRotor Original Position:", self.origPos)

    #take a character as input and return its rotated ciphertext character
    #rotate rotor after encrypting automatically
    def outputRotated(self, charX):
        #Rotate the rotor if the flag is set to true
        if(self.rotateFlag):
            self.rotate()

        X_index = Char_Indices[charX]
        rotatedChar = self.rotorPermutation[X_index]

        return rotatedChar

    def outputRotatedRev(self, charX):
        X_index = getIndex(charX, self.rotorPermutation)
        rotatedCharRev = Letters[X_index]

        return rotatedCharRev

#Class to Implement a Set of Rotors along with the reflector
#Customized to build a 3-Rotor Set and Fastest moving rotor is closest to Plugboard
#Both the parameters in the init uniquely identify the rotor element of the key

class RotorSet(Rotor, Reflector):
    'Implements a Rotor Set for a 3-Rotor Enigma Machine'
    #useRotors is a list that has the index of the rotor to be used from Enigma_Rotors
    #(Rotor1 = 0, Rotor2 = 1, ..., Rotor8 = 7)
    #rotorPositions is list with the initial position of the rotors being used
    #both the above lists are in order of closest to plugboard to furthest in terms of their rotor set position

    def __init__(self, useRotors, rotorPositions, reflectNum): #Instantiates the rotor set
        self.R1 = Rotor(Enigma_Rotors[useRotors[0]], rotorPositions[0]) #Rotor R1 to be placed nearest the plugboard
        self.R2 = Rotor(Enigma_Rotors[useRotors[1]], rotorPositions[1])
        self.R3 = Rotor(Enigma_Rotors[useRotors[2]], rotorPositions[2])

        #Instantiate correct reflector
        self.Reflector1 = Reflector(reflectNum)

        print("\nInitialized Enigma 3-Rotor Set with Plugboard Successfully")
        #self.R1.displayRotor()
        #self.R2.displayRotor()
        #self.R3.displayRotor()
        print("-------------------------------------------------------------------------------------------------------")

    #funtion to display the state of the rotor positions
    def displayRotorState(self):
        print("Current RotorState:", self.R1.rotorPos, self.R2.rotorPos, self.R3.rotorPos)

    #Encrypt text as per the rotor settings and spin rotors as per the gearing
    #Each rotor takes the output of the previous rotor as its input
    #The reflector and backward path has also been accomodated
    def encryptText(self, textArray):
        rotorSet_out = []
        #Only set rotate flag for R1=True since it must rotate every turn
        self.R1.rotateFlag = True
        #Traverse the array with characters once
        for indx, letter in enumerate(textArray):

            #self.displayRotorState() #Check Rotorstate
            #R1 always encrypts the letter and rotates
            R1out = self.R1.outputRotated(letter)

            '''
            #Test just one rotor + reflector
            refOut = self.Reflector1.outputReflected(R1out)

            R1outrev = self.R1.outputRotatedRev(refOut)

            print("Letter Path:", letter, " ->", R1out, " ->", refOut, " ->"
                 , R1outrev)
            '''
            #R2 is configured to rotate every time R1 completes a cycle, except first index
            if(self.R1.rotorPos == ((self.R1.origPos)%26) and indx != 0 ):
                self.R2.rotateFlag = True

            R2out = self.R2.outputRotated(R1out)

            #R3 is configured to rotate every time R3 completes a cycle, except first cycle
            if(self.R2.rotorPos == ((self.R2.origPos)%26) and self.R2.rotateFlag and indx > 26):
                self.R3.rotateFlag = True

            R3out = self.R3.outputRotated(R2out)

            #Pass through reflector
            reflectedOut = self.Reflector1.outputReflected(R3out)

            #Pass through R3 again
            R3out_rev = self.R3.outputRotatedRev(reflectedOut)
            self.R3.rotateFlag = False #Rotate flag must stay false as default

            #Pass through R2 again
            R2out_rev = self.R2.outputRotatedRev(R3out_rev)
            self.R2.rotateFlag = False #Rotate flag must be set to false as default

            R1out_rev = self.R1.outputRotatedRev(R2out_rev)

            #print path of letter:
            #print("Letter Path:", letter, " ->", R1out, " ->", R2out, " ->", R3out, " ->", reflectedOut, " ->"
                 #, R3out_rev, " ->", R2out_rev, " ->", R1out_rev)
            #The character stored in R1out_rev is the final result of rotors + reflector

            rotorSet_out.append(R1out_rev)

        return rotorSet_out

max_stars_repo_name: kvombatkere/Enigma-Machine

max_issues_repo_name: kvombatkere/Enigma-Machine

max_forks_repo_name: kvombatkere/Enigma-Machine


text:
import numpy as np

class Document:

    def __init__(self, file_path, encoding):
        """
        Stores information about the given file.
        :param file_path: str, path to the file name to be read in.
        :param encoding: the encoding to be used for the file.
        """
        self.file_path = file_path
        self.encoding = encoding
        self.contents = None

    def store_contents(self):
        """
        Reads in the file's contents and stores it in self.contents.
        :return: None.
        """
        contents = list()
        with open(self.file_path, "r", encoding=self.encoding) as f:
            pos = 0
            for line in f:
                try:
                    sentences = line.split(".")
                    for sentence in sentences:
                        if sentence and sentence != ['\n']:
                            contents.append((sentence, pos))
                            pos += 1
                except UnicodeDecodeError:
                    pass
        self.contents = contents

    def get_embeddings(self):
        """
        Only to be called on the word embeddings file.
        Extracts the word embeddings and stores them as a dictionary in self.contents.
        :return: None.
        """
        word_embeddings = dict()

        f = open(self.file_path, encoding=self.encoding)  # GloVe word embeddings

        for line in f:
            values = line.split()
            word = values[0]
            coefs = np.asarray(values[1:], dtype="float32")
            word_embeddings[word] = coefs
        f.close()

        self.contents = word_embeddings

max_stars_repo_name: aleph-aseffa/compress

max_issues_repo_name: aleph-aseffa/compress

max_forks_repo_name: aleph-aseffa/compress


text:
# Copyright Hugh Perkins 2016
"""
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.

alloca is created before we know which address space it should be
we need to handle this, but first we need a test case for this
"""
import numpy as np
import pyopencl as cl
import os
import subprocess
from test import test_common
from test.test_common import offset_type

basename = os.path.basename(__name__).split('.')[-1]
print('basename', basename)

def test_alloca(context, q, float_data, float_data_gpu):

    code = """
__global__ void mykernel(float *data) {
    float *foo = data;
    foo[0] = data[1] + data[2] + data[3];
}
"""
    kernelName = test_common.mangle('mykernel', ['float *'])
    try:
        dict = test_common.compile_code_v2(cl, context, code, kernelName, num_clmems=1)
        prog = dict['prog']
        cl_sourcecode = dict['cl_sourcecode']
        print('cl_sourcecode', cl_sourcecode)
    except Exception as e:
        with open('/tmp/testprog-device.cl', 'r') as f:
            print(f.read())
        raise e
    print('type(offset_type(0))', type(offset_type(0)))
    prog.__getattr__(kernelName)(
        q, (32,), (32,),
        float_data_gpu, offset_type(0), offset_type(0), cl.LocalMemory(4))
    # q.finish()
    # cl.enqueue_copy(q, float_data, float_data_gpu)
    # q.finish()

max_stars_repo_name: dendisuhubdy/coriander

max_issues_repo_name: hughperkins/cuda-ir-to-opencl

max_forks_repo_name: hughperkins/cuda-ir-to-opencl