Source code for schrodinger.application.matsci.stress_strain

"""
Utilities for stress strain.

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

import os
from collections import defaultdict

import numpy

from schrodinger.application.matsci import jobutils
from schrodinger.application.matsci import msprops

AXIS_HEADER = 'Axis %s'
AXIS_HEADERS = [AXIS_HEADER % axis for axis in 'abc']
AXIS_AVG = 'Average'

STRESS_EFF_XDATA = {'Effective': (0, 5), 'Normal': (1, 5), 'Transverse': (3, 5)}

RESIDUAL_STRAIN_Y_COL = 16

# Without and with residual strain
CSV_DATA_NCOLS = [16, 17]


def get_CSV_path(st):
    """
    Return the CSV path from the structure.

    :type st: schrodinger.structure.Structure
    :param st: the structure

    :rtype: str or None
    :return: the CSV path or None if one can not be found
    """

    cms_fn = st.property.get(msprops.ORIGINAL_CMS_PROP, '')
    if cms_fn.endswith('-out.cms'):
        csv_fn = cms_fn.replace('-out.cms', '.csv')
        source_path = jobutils.get_source_path(st)
        csv_fn = os.path.join(source_path, csv_fn)
        if os.path.exists(csv_fn):
            return csv_fn


def parse_CSV(csv_fn):
    """
    Parse CSV file from the calculation.

    :type csv_fn: str
    :param csv_fn: Path to CSV file

    :rtype: dict, dict
    :return: Both dictionaries keys are main axis name. First dictionary with
        values is the associated data for the axis and second dictionary
        the header string for the axis.

    :raise ValueError: If parsing of CSV file fails
    """

    data = {}
    headers = defaultdict(str)
    start_reading = False

    with open(csv_fn) as csv_fh:
        for line in csv_fh:
            if 'Main axis =' in line:
                tmp = line.split(',')
                direction = AXIS_HEADER % tmp[1].strip()
                if direction not in AXIS_HEADERS:
                    raise ValueError('Unknown axis: %s' % direction)

                start_reading = True

            if start_reading:
                headers[direction] += line

            if 'eps_eff' in line:
                while True:
                    line = csv_fh.readline().strip()
                    if not line:
                        break
                    tmp = list(map(float, line.split(',')))
                    if len(tmp) not in CSV_DATA_NCOLS:
                        raise ValueError('Different CSV format.')

                    if direction not in data:
                        data[direction] = numpy.array([tmp])
                    else:
                        data[direction] = numpy.append(data[direction], [tmp],
                                                       axis=0)

    if not data:
        raise ValueError('Could not parse CSV file.')

    nkeys = len(data.keys())
    if nkeys == 1:
        return data, headers

    avg_data = None
    for key, val in data.items():
        if avg_data is None:
            avg_data = numpy.array(val)
        else:
            avg_data += numpy.array(val)

    avg_data /= nkeys
    data[AXIS_AVG] = avg_data
    return data, headers