#!/usr/bin/env python
#
# Copyright 2019 DFKI GmbH.
#
# Permission is hereby granted, free of charge, to any person obtaining a
# copy of this software and associated documentation files (the
# "Software"), to deal in the Software without restriction, including
# without limitation the rights to use, copy, modify, merge, publish,
# distribute, sublicense, and/or sell copies of the Software, and to permit
# persons to whom the Software is furnished to do so, subject to the
# following conditions:
#
# The above copyright notice and this permission notice shall be included
# in all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS
# OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
# MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN
# NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM,
# DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR
# OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE
# USE OR OTHER DEALINGS IN THE SOFTWARE.
"""
Functions for retargeting based on the paper "Using an Intermediate Skeleton and Inverse Kinematics for Motion Retargeting"
by Monzani et al.
See: http://www.vis.uni-stuttgart.de/plain/vdl/vdl_upload/91_35_retargeting%20monzani00using.pdf
"""
import numpy as np
import math
from .constants import OPENGL_UP_AXIS, GAME_ENGINE_SPINE_OFFSET_LIST
from transformations import quaternion_matrix, quaternion_multiply, quaternion_about_axis, quaternion_inverse, quaternion_from_matrix
from .utils import normalize, align_axis, find_rotation_between_vectors, align_root_translation, to_local_cos, get_quaternion_rotation_by_name, apply_additional_rotation_on_frames, project_vector_on_axis, quaternion_from_vector_to_vector
from ..animation_data.skeleton_models import JOINT_CHILD_MAP
from .analytical import create_local_cos_map_from_skeleton_axes_with_map

JOINT_CHILD_MAP = dict()
JOINT_CHILD_MAP["root"] = "pelvis"
JOINT_CHILD_MAP["pelvis"] = "spine_2"
JOINT_CHILD_MAP["spine_2"] = "neck"
JOINT_CHILD_MAP["neck"] = "head"
JOINT_CHILD_MAP["left_clavicle"] = "left_shoulder"
JOINT_CHILD_MAP["left_shoulder"] = "left_elbow"
JOINT_CHILD_MAP["left_elbow"] = "left_wrist"
JOINT_CHILD_MAP["left_wrist"] = "left_finger"
JOINT_CHILD_MAP["right_clavicle"] = "right_shoulder"
JOINT_CHILD_MAP["right_shoulder"] = "right_elbow"
JOINT_CHILD_MAP["right_elbow"] = "right_wrist"
JOINT_CHILD_MAP["right_wrist"] = "right_finger"
JOINT_CHILD_MAP["left_hip"] = "left_knee"
JOINT_CHILD_MAP["left_knee"] = "left_ankle"
JOINT_CHILD_MAP["right_elbow"] = "right_wrist"
JOINT_CHILD_MAP["right_hip"] = "right_knee"
JOINT_CHILD_MAP["right_knee"] = "right_ankle"
JOINT_CHILD_MAP["left_ankle"] = "left_toe"
JOINT_CHILD_MAP["right_ankle"] = "right_toe"

def estimate_correction(target_zero_vector_y, target_zero_vector_x,  src_zero_vector_y, src_zero_vector_x):
    q = quaternion_from_vector_to_vector(target_zero_vector_y, src_zero_vector_y)
    q = normalize(q)
    m = quaternion_matrix(q)[:3, :3]
    target_zero_vector_x = normalize(np.dot(m, target_zero_vector_x))
    qx = quaternion_from_vector_to_vector(target_zero_vector_x, src_zero_vector_x)
    q = quaternion_multiply(qx, q)
    q = normalize(q)
    return q

def create_correction_map(target_skeleton, target_to_src_joint_map, target_cos_map, src_cos_map):
    correction_map = dict()
    joint_map = target_skeleton.skeleton_model["joints"]
    for target_name in target_to_src_joint_map:
        src_name = target_to_src_joint_map[target_name]
        if src_name in src_cos_map and target_name is not None:
            src_zero_vector_y = src_cos_map[src_name]["y"]
            target_zero_vector_y = target_cos_map[target_name]["y"]
            src_zero_vector_x = src_cos_map[src_name]["x"]
            target_zero_vector_x = target_cos_map[target_name]["x"]
            if target_zero_vector_y is not None and src_zero_vector_y is not None:
                q = estimate_correction(target_zero_vector_y, target_zero_vector_x,  src_zero_vector_y, src_zero_vector_x)
                correction_map[target_name] = q
    return correction_map

def get_quaternion_to_axis(skeleton, joint_a, joint_b, axis):
    ident_f = skeleton.identity_frame
    ap = skeleton.nodes[joint_a].get_global_position(ident_f)
    bp = skeleton.nodes[joint_b].get_global_position(ident_f)
    delta = bp - ap
    delta /= np.linalg.norm(delta)
    return quaternion_from_vector_to_vector(axis, delta)

def rotate_axes2(cos, q):
    m = quaternion_matrix(q)[:3, :3]
    aligned_axes = dict()
    for key, a in list(cos.items()):
        aligned_axes[key] = np.dot(m, a)
        aligned_axes[key] = normalize(aligned_axes[key])
    return aligned_axes

def get_child_joint(skeleton_model, inv_joint_map, node_name, src_children_map):
    """ Warning output is random if there are more than one child joints
        and the value is not specified in the JOINT_CHILD_MAP """
    child_name = None
    if node_name in src_children_map and len(src_children_map[node_name]) > 0:
        child_name = src_children_map[node_name][-1]
    if node_name in inv_joint_map:
        joint_name = inv_joint_map[node_name]
        while joint_name in JOINT_CHILD_MAP:
            _child_joint_name = JOINT_CHILD_MAP[joint_name]

            # check if child joint is mapped
            joint_key = None
            if _child_joint_name in skeleton_model["joints"]:
                joint_key = skeleton_model["joints"][_child_joint_name]

            if joint_key is not None: # return child joint
                child_name = joint_key
                return child_name
            else: #keep traversing until end of child map is reached
                if _child_joint_name in JOINT_CHILD_MAP:
                    joint_name = JOINT_CHILD_MAP[_child_joint_name]
                    #print(joint_name)
                else:
                    break
    return child_name

def rotate_axes_in_place(cos, q):
    m = quaternion_matrix(q)[:3, :3]
    for key, a in list(cos.items()):
        cos[key] = np.dot(m, a)
        cos[key] = normalize(cos[key])
    return cos

def align_axis_in_place(axes, key, new_vec):
    q = quaternion_from_vector_to_vector(axes[key], new_vec)
    aligned_axes = rotate_axes_in_place(axes, q)
    return q, aligned_axes

def to_local_cos_fast(skeleton, node_name, frame, q):
    # bring into parent coordinate system
    pm = np.array(skeleton.nodes[node_name].get_global_matrix(frame, use_cache=True)[:3,:3])
    inv_p = quaternion_inverse(quaternion_from_matrix(pm))
    return quaternion_multiply(inv_p, q)

def align_root_joint(new_skeleton, free_joint_name, axes, global_src_up_vec, global_src_x_vec,joint_cos_map, max_iter_count=10):
    # handle special case for the root joint
    # apply only the y axis rotation of the Hip to the Game_engine node
    q = [1, 0, 0, 0]
    #apply first time
    qx, axes = align_axis_in_place(axes, "x", global_src_x_vec)  # first find rotation to align x axis
    q = quaternion_multiply(qx, q)
    q = normalize(q)

    qy, axes = align_axis_in_place(axes, "y", global_src_up_vec)  # then add a rotation to let the y axis point up
    q = quaternion_multiply(qy, q)
    q = normalize(q)

    #apply second time
    qx, axes = align_axis_in_place(axes, "x", global_src_x_vec)  # first find rotation to align x axis
    q = quaternion_multiply(qx, q)
    q = normalize(q)
    qy, axes = align_axis_in_place(axes, "y", global_src_up_vec)  # then add a rotation to let the y axis point up
    q = quaternion_multiply(qy, q)
    q = normalize(q)

    # print("handle special case for pelvis")
    # handle special case of applying the x axis rotation of the Hip to the pelvis
    node = new_skeleton.nodes[free_joint_name]
    t_pose_global_m = node.get_global_matrix(new_skeleton.reference_frame)[:3, :3]
    global_original = np.dot(t_pose_global_m, joint_cos_map[free_joint_name]["y"])
    global_original = normalize(global_original)
    qoffset = find_rotation_between_vectors(OPENGL_UP_AXIS, global_original)
    q = quaternion_multiply(q, qoffset)
    q = normalize(q)
    return q

def align_joint(local_target_axes, up_vec, x_vec):
    q = [1, 0, 0, 0]
    qy, axes = align_axis_in_place(local_target_axes, "y", up_vec)
    q = quaternion_multiply(qy, q)
    q = normalize(q)

    # then align the twisting angles
    if x_vec is not None:
        qx, axes = align_axis_in_place(axes, "x", x_vec)
        q = quaternion_multiply(qx, q)
        q = normalize(q)
        # print("set twist angle", free_joint_name, twist_angle)
    return q

def find_rotation_analytically(new_skeleton, free_joint_name, target, frame, joint_cos_map, is_root=False, max_iter_count=10, twist_angle=None):
    global_src_up_vec = target[0]
    if twist_angle is None:
        global_src_x_vec = target[1]
    else:
        global_src_x_vec = None
    local_target_axes = dict(joint_cos_map[free_joint_name])

    if is_root:
        q = align_root_joint(new_skeleton, free_joint_name, local_target_axes, global_src_up_vec,global_src_x_vec, joint_cos_map, max_iter_count)

    else:
        # first align the bone vectors
        q = [1, 0, 0, 0]
        qy, axes = align_axis_in_place(local_target_axes, "y", global_src_up_vec)
        q = quaternion_multiply(qy, q)
        q = normalize(q)

        # then align the twisting angles
        if global_src_x_vec is not None:
            qx, axes = align_axis_in_place(axes, "x", global_src_x_vec)
            q = quaternion_multiply(qx, q)
            q = normalize(q)

    #if "FK" in free_joint_name:
    #    q = to_local_cos(new_skeleton, free_joint_name, frame, q)
    if new_skeleton.nodes[free_joint_name].parent is not None:
        #if "upLeg" in free_joint_name: # it does not work for the legs for some reason
        #    q = to_local_cos(new_skeleton, new_skeleton.nodes[free_joint_name].parent.node_name, frame, q)
        #else:
        q = to_local_cos(new_skeleton, new_skeleton.nodes[free_joint_name].parent.node_name, frame, q)

    if twist_angle is not None:
        # separate rotation
        local_twist_axis = np.array(joint_cos_map[free_joint_name]["y"])
        swing_q, twist_q = swing_twist_decomposition(q, local_twist_axis)
        # replace
        twist_q = quaternion_about_axis(-twist_angle, local_twist_axis)
        q = quaternion_multiply(swing_q, twist_q)
        q = normalize(q)
    return q

def find_rotation_analytically_with_guess(new_skeleton, free_joint_name, target, frame, joint_cos_map, prev_global_q, is_root=False, max_iter_count = 10):
    global_src_up_vec = target[0]
    global_src_x_vec = target[1]
    local_target_axes = joint_cos_map[free_joint_name]
    rotated_axes = rotate_axes2(local_target_axes, prev_global_q)
    #print("rotate",local_target_axes, rotated_axes, prev_global_q)
    #print("")
    if is_root:
        q = align_root_joint(new_skeleton, free_joint_name, rotated_axes, global_src_up_vec,global_src_x_vec, joint_cos_map, max_iter_count)
    else:
        q = align_joint(rotated_axes, global_src_up_vec, global_src_x_vec)
    q = quaternion_multiply(q, prev_global_q)
    q = normalize(q)
    return to_local_cos_fast(new_skeleton, free_joint_name, frame, q)

def get_parent_map(joints):
    """Returns a dict of node names to their parent node's name"""
    parent_dict = dict()
    for joint_name in list(joints.keys()):
        parent_dict[joint_name] = joints[joint_name]['parent']
    return parent_dict

def get_children_map(joints):
    """Returns a dict of node names to a list of children names"""
    child_dict = dict()
    for joint_name in list(joints.keys()):
        parent_name = joints[joint_name]['parent']
        if parent_name not in child_dict:
            child_dict[parent_name] = list()
        child_dict[parent_name].append(joint_name)
    return child_dict

def swing_twist_decomposition(q, twist_axis):
    """ code by janis sprenger based on
        Dobrowsolski 2015 Swing-twist decomposition in Clifford algebra. https://arxiv.org/abs/1506.05481
    """
    #q = normalize(q)
    #twist_axis = np.array((q * offset))[0]
    projection = np.dot(twist_axis, np.array([q[1], q[2], q[3]])) * twist_axis
    twist_q = np.array([q[0], projection[0], projection[1],projection[2]])
    if np.linalg.norm(twist_q) == 0:
        twist_q = np.array([1,0,0,0])
    twist_q = normalize(twist_q)
    swing_q = quaternion_multiply(q, quaternion_inverse(twist_q))#q * quaternion_inverse(twist)
    return swing_q, twist_q

class PointCloudRetargeting(object):
    def __init__(self, src_joints, src_model, target_skeleton, target_to_src_joint_map, scale_factor=1.0, additional_rotation_map=None, constant_offset=None, place_on_ground=False, ground_height=0):
        self.src_joints = src_joints
        self.src_model = src_model
        self.target_skeleton = target_skeleton
        self.target_to_src_joint_map = target_to_src_joint_map
        if target_skeleton.skeleton_model["joints"]["pelvis"] is not None:
            self.target_skeleton_root = target_skeleton.skeleton_model["joints"]["pelvis"]
        else:
            self.target_skeleton_root = target_skeleton.root

        #FIXME: enable spine during retargeting
        for j in [ "spine_1", "spine"]:#"spine_2",
            k = self.target_skeleton.skeleton_model["joints"][j]
            self.target_to_src_joint_map[k] = None

        self.src_to_target_joint_map = {v: k for k, v in list(self.target_to_src_joint_map.items())}
        self.scale_factor = scale_factor
        self.n_params = len(self.target_skeleton.animated_joints) * 4 + 3
        self.ground_height = ground_height
        self.additional_rotation_map = additional_rotation_map
        self.src_inv_joint_map = dict((v,k) for k, v in src_model["joints"].items())
        self.src_child_map = dict()
        self.src_parent_map = get_parent_map(src_joints)
        src_children_map = get_children_map(src_joints)
        for src_name in self.src_joints:
            src_child = get_child_joint(self.src_model, self.src_inv_joint_map, src_name, src_children_map)
            if src_child is not None:
                self.src_parent_map[src_child] = src_name
                self.src_child_map[src_name] = src_child
            else:
                self.src_child_map[src_name] = None
        #print("ch",self.src_child_map)
        #for j in ["pelvis", "spine", "spine_1", "spine_2"]:
        #    if j in target_joints:
        src_joint_map = self.src_model["joints"]
        for j in ["neck", "spine_2", "spine_1", "spine"]:
            if j in src_joint_map:
                self.src_parent_map["spine_03"] = "pelvis"
        self.src_child_map[src_joint_map["pelvis"]] = src_joint_map["neck"]#"pelvis" "neck_01"

        self.constant_offset = constant_offset
        self.place_on_ground = place_on_ground
        self.temp_frame_data = dict()

        self.target_cos_map = create_local_cos_map_from_skeleton_axes_with_map(self.target_skeleton)
        if "cos_map" in target_skeleton.skeleton_model:
            self.target_cos_map.update(target_skeleton.skeleton_model["cos_map"])

        target_joints = self.target_skeleton.skeleton_model["joints"]
        self.target_spine_joints = [target_joints[j] for j in ["neck", "spine_2", "spine_1", "spine"] if j in target_joints]#["spine_03", "neck_01"]
        self.target_ball_joints = [target_joints[j] for j in ["left_shoulder", "right_shoulder", "left_hip", "right_hip"] if j in target_joints]# ["thigh_r", "thigh_l", "upperarm_r", "upperarm_l"]
        self.target_ankle_joints = [target_joints[j] for j in ["left_ankle", "right_ankle"] if j in target_joints]
        self.clavicle_joints = [target_joints[j] for j in ["right_clavicle", "left_clavicle"] if j in target_joints]
        self.twist_angle_joints = [target_joints[j] for j in ["right_knee","left_knee","left_shoulder", "right_shoulder", "right_clavicle", "left_clavicle","right_elbow", "left_elbow","left_wrist", "right_wrist"] if j in target_joints]
        if "neck" in target_joints:
            self.target_neck_joint = target_joints["neck"]
        else:
            self.target_neck_joint = None

        left_hip = self.src_model["joints"]["left_hip"]
        right_hip = self.src_model["joints"]["right_hip"]
        self.left_hip_idx = self.src_joints[left_hip]["index"]
        self.right_hip_idx = self.src_joints[right_hip]["index"]
        left_shoulder = self.src_model["joints"]["left_shoulder"]
        right_shoulder = self.src_model["joints"]["right_shoulder"]
        self.left_shoulder_idx = self.src_joints[left_shoulder]["index"]
        self.right_shoulder_idx = self.src_joints[right_shoulder]["index"]
        self.ref_rotation = dict()
        for target_name in self.target_skeleton.animated_joints:
            self.ref_rotation[target_name] = get_quaternion_rotation_by_name(target_name, self.target_skeleton.reference_frame,
                                                                             self.target_skeleton, root_offset=3)

    def estimate_src_joint_cos(self, src_name, child_name, target_name, src_frame):
        joint_idx = self.src_joints[src_name]["index"]
        child_idx = self.src_joints[child_name]["index"]
        global_src_up_vec = src_frame[child_idx] - src_frame[joint_idx]
        global_src_up_vec /= np.linalg.norm(global_src_up_vec)
        self.temp_frame_data[src_name] = global_src_up_vec
        if target_name == self.target_skeleton.skeleton_model["joints"]["pelvis"]:
            global_src_x_vec = src_frame[self.left_hip_idx] - src_frame[self.right_hip_idx]
            global_src_x_vec /= np.linalg.norm(global_src_x_vec)
        elif target_name in self.target_spine_joints or target_name == "CC_Base_Waist":  # find x vector from shoulders
            global_src_x_vec = src_frame[self.left_shoulder_idx] - src_frame[self.right_shoulder_idx]
            global_src_x_vec /= np.linalg.norm(global_src_x_vec)
        elif target_name in self.target_ball_joints:  # use x vector of child
            child_child_name = self.src_child_map[child_name]
            child_child_idx = self.src_joints[child_child_name]["index"]
            child_global_src_up_vec = src_frame[child_child_idx] - src_frame[child_idx]
            child_global_src_up_vec /= np.linalg.norm(child_global_src_up_vec)
            global_src_x_vec = np.cross(global_src_up_vec, child_global_src_up_vec)
            global_src_x_vec /= np.linalg.norm(global_src_x_vec)
        else:  # find x vector by cross product with parent
            global_src_x_vec = None
            if src_name in self.src_parent_map:
                parent_joint = self.src_parent_map[src_name]
                if parent_joint in self.temp_frame_data:
                    global_parent_up_vector = self.temp_frame_data[parent_joint]
                    global_src_x_vec = np.cross(global_src_up_vec, global_parent_up_vector)
                    global_src_x_vec /= np.linalg.norm(global_src_x_vec)
                    if target_name not in self.target_ankle_joints:
                        global_src_x_vec = -global_src_x_vec
        return global_src_up_vec, global_src_x_vec

    def rotate_bone(self, src_name, target_name, src_frame, target_frame, guess, pose_angles=None):
        q = guess
        if src_name not in self.src_child_map.keys() or self.src_child_map[src_name] is None:
            return q
        if self.src_child_map[src_name] in self.src_to_target_joint_map:#  and or target_name =="neck_01" or target_name.startswith("hand")
            child_name = self.src_child_map[src_name]
            if child_name not in self.src_joints.keys():
                return q
            is_root = False
            if target_name == self.target_skeleton_root:
                is_root = True
            src_cos = self.estimate_src_joint_cos(src_name, child_name, target_name, src_frame)
            if src_cos[1] is None:
                return q
            twist_angle = None
            if pose_angles is not None and self.target_skeleton.nodes[target_name].parent is not None and target_name in self.twist_angle_joints:
                joint_idx = self.src_joints[src_name]["index"]
                twist_angle = pose_angles[joint_idx][0]
            q = find_rotation_analytically(self.target_skeleton, target_name, src_cos, target_frame, self.target_cos_map, is_root=is_root, twist_angle=twist_angle)
        return q/np.linalg.norm(q)

    def retarget_frame(self, src_frame, ref_frame, pose_angles=None):

        self.target_skeleton.clear_cached_global_matrices()
        target_frame = np.zeros(self.n_params)
        self.temp_frame_data.clear()
        # copy the root translation assuming the rocketbox skeleton with static offset on the hips is used as source
        target_frame[:3] = np.array(src_frame[0]) * self.scale_factor
        if self.constant_offset is not None:
            target_frame[:3] += self.constant_offset
        animated_joints = self.target_skeleton.animated_joints
        target_offset = 3
        for target_name in animated_joints:
            q = self.ref_rotation[target_name]
            if target_name in self.target_to_src_joint_map.keys():
                src_name = self.target_to_src_joint_map[target_name]
                if src_name is not None and src_name in self.src_joints.keys():
                    q = self.rotate_bone(src_name, target_name, src_frame, target_frame, q, pose_angles)
            #if ref_frame is not None:
            #    q = q if np.dot(ref_frame[target_offset:target_offset + 4], q) >= 0 else -q
            target_frame[target_offset:target_offset + 4] = q
            target_offset += 4
        return target_frame

    def run(self, src_frames, frame_range):
        n_frames = len(src_frames)
        target_frames = []
        if n_frames > 0:
            if frame_range is None:
                frame_range = (0, n_frames)
            if self.additional_rotation_map is not None:
               src_frames = apply_additional_rotation_on_frames(self.src_skeleton.animated_joints, src_frames, self.additional_rotation_map)

            ref_frame = None
            for idx, src_frame in enumerate(src_frames[frame_range[0]:frame_range[1]]):
                target_frame = self.retarget_frame(src_frame, ref_frame)
                if ref_frame is None:
                    ref_frame = target_frame
                target_frames.append(target_frame)
            target_frames = np.array(target_frames)
            if self.place_on_ground:
                delta = target_frames[0][1] - self.ground_height
                target_frames[:, 1] -= delta
        return target_frames

def generate_joint_map(src_model, target_model):
    joint_map = dict()
    for j in src_model["joints"]:
        if j in target_model["joints"]:
            src = src_model["joints"][j]
            target = target_model["joints"][j]
            joint_map[target] = src
    return joint_map

def retarget_from_point_cloud_to_target(src_joints, src_model, target_skeleton, src_frames, joint_map=None, additional_rotation_map=None, scale_factor=1.0, frame_range=None, place_on_ground=False):
    if joint_map is None:
        joint_map = generate_joint_map(src_model, target_skeleton.skeleton_model)
    retargeting = PointCloudRetargeting(src_joints, src_model, target_skeleton, joint_map, scale_factor, additional_rotation_map=additional_rotation_map, place_on_ground=place_on_ground)
    return retargeting.run(src_frames, frame_range)


#!/usr/bin/env python

"""
This module provides classes to create phase diagrams.
"""

from __future__ import division

__author__ = "Shyue Ping Ong"
__copyright__ = "Copyright 2011, The Materials Project"
__version__ = "2.0"
__maintainer__ = "Shyue Ping Ong"
__email__ = "shyue@mit.edu"
__status__ = "Production"
__date__ = "Nov 25, 2012"

import collections
import logging

import numpy as np

from pyhull.convex_hull import ConvexHull

from pymatgen.core.composition import Composition
from pymatgen.phasediagram.entries import GrandPotPDEntry, TransformedPDEntry

from pymatgen.core.periodic_table import DummySpecie
from pymatgen.analysis.reaction_calculator import Reaction, ReactionError

logger = logging.getLogger(__name__)

class PhaseDiagram (object):
    """
    Simple phase diagram class taking in elements and entries as inputs.
    The algorithm is based on the work in the following papers:

    1. S. P. Ong, L. Wang, B. Kang, and G. Ceder, Li-Fe-P-O2 Phase Diagram from
       First Principles Calculations. Chem. Mater., 2008, 20(5), 1798-1807.
       doi:10.1021/cm702327g

    2. S. P. Ong, A. Jain, G. Hautier, B. Kang, G. Ceder, Thermal stabilities
       of delithiated olivine MPO4 (M=Fe, Mn) cathodes investigated using first
       principles calculations. Electrochem. Comm., 2010, 12(3), 427-430.
       doi:10.1016/j.elecom.2010.01.010

    .. attribute: elements:

        Elements in the phase diagram.

    ..attribute: all_entries

        All entries provided for Phase Diagram construction. Note that this
        does not mean that all these entries are actually used in the phase
        diagram. For example, this includes the positive formation energy
        entries that are filtered out before Phase Diagram construction.

    .. attribute: qhull_data

        Data used in the convex hull operation. This is essentially a matrix of
        composition data and energy per atom values created from qhull_entries.

    .. attribute: dim

        The dimensionality of the phase diagram.

    .. attribute: facets

        Facets of the phase diagram in the form of  [[1,2,3],[4,5,6]...]

    .. attribute: el_refs:

        List of elemental references for the phase diagrams. These are
        entries corresponding to the lowest energy element entries for simple
        compositional phase diagrams.

    .. attribute: qhull_entries:

        Actual entries used in convex hull. Excludes all positive formation
        energy entries.
    """

    # Tolerance for determining if formation energy is positive.
    formation_energy_tol = 1e-11

    def __init__(self, entries, elements=None):
        """
        Standard constructor for phase diagram.

        Args:
            entries:
                A list of PDEntry-like objects having an energy,
                energy_per_atom and composition.
            elements:
                Optional list of elements in the phase diagram. If set to None,
                the elements are determined from the the entries themselves.
        """
        if elements is None:
            elements = set()
            map(elements.update, [entry.composition.elements
                                  for entry in entries])
        self.all_entries = entries
        self.elements = tuple(elements)
        dim = len(self.elements)
        self.el_refs = self._get_el_refs()
        (self.qhull_entries, qhull_data) = self._create_convhull_data()
        if len(qhull_data) == dim:
            self.facets = [range(dim)]
        else:
            facets = ConvexHull(qhull_data).vertices
            logger.debug("Final facets are\n{}".format(facets))

            logger.debug("Removing vertical facets...")
            finalfacets = []
            for facet in facets:
                facetmatrix = np.zeros((len(facet), len(facet)))
                count = 0
                is_element_facet = True
                for vertex in facet:
                    facetmatrix[count] = np.array(qhull_data[vertex])
                    facetmatrix[count, dim - 1] = 1
                    count += 1
                    if len(self.qhull_entries[vertex].composition) > 1:
                        is_element_facet = False
                if abs(np.linalg.det(facetmatrix)) > 1e-8 and\
                        (not is_element_facet):
                    finalfacets.append(facet)
                else:
                    logger.debug("Removing vertical facet : {}".format(facet))
            self.facets = finalfacets
        self.qhull_data = qhull_data
        self.dim = dim

    @property
    def unstable_entries(self):
        """
        Entries that are unstable in the phase diagram. Includes positive
        formation energy entries.
        """
        return [e for e in self.all_entries if e not in self.stable_entries]

    @property
    def stable_entries(self):
        """
        Returns the stable entries in the phase diagram.
        """
        stable_entries = set()
        for facet in self.facets:
            for vertex in facet:
                stable_entries.add(self.qhull_entries[vertex])
        return stable_entries

    @property
    def all_entries_hulldata(self):
        """
        Same as qhull_data, but for all entries rather than just negative
        formation energy ones.
        """
        return self._process_entries_qhulldata(self.all_entries)

    def get_form_energy(self, entry):
        """
        Returns the formation energy for an entry (NOT normalized) from the
        elemental references.

        Args:
            entry:
                A PDEntry-like object.

        Returns:
            Formation energy from the elemental references.
        """
        comp = entry.composition
        energy = entry.energy - sum([comp[el] *
                                     self.el_refs[el].energy_per_atom
                                     for el in comp.elements])
        return energy

    def get_form_energy_per_atom(self, entry):
        """
        Returns the formation energy per atom for an entry from the
        elemental references.

        Args:
            entry:
                An PDEntry-like object

        Returns:
            Formation energy **per atom** from the elemental references.
        """
        comp = entry.composition
        return self.get_form_energy(entry) / comp.num_atoms

    def _process_entries_qhulldata(self, entries_to_process):
        """
        From a sequence of entries, generate the necessary for the convex hull.
        Using the Li-Fe-O phase diagram as an example, this is of the form:
        [[ Fe_fraction_entry_1, O_fraction_entry_1, Energy_per_atom_entry_1],
         [ Fe_fraction_entry_2, O_fraction_entry_2, Energy_per_atom_entry_2],
         ...]]

        Note that there are only two independent variables, since the third
        elemental fraction is fixed by the constraint that all compositions sum
        to 1. The choice of the elements is arbitrary.
        """
        def make_row(entry):
            comp = entry.composition
            row = [comp.get_atomic_fraction(self.elements[i])
                   for i in xrange(1, len(self.elements))]
            row.append(entry.energy_per_atom)
            return row

        return map(make_row, entries_to_process)

    def _get_el_refs(self):
        el_refs = {}
        for el in self.elements:
            el_entries = filter(lambda e: e.composition.is_element and
                                e.composition.elements[0] == el,
                                self.all_entries)
            if len(el_entries) == 0:
                raise PhaseDiagramError("There are no entries associated with"
                                        " terminal {}.".format(el))
            el_refs[el] = min(el_entries, key=lambda e: e.energy_per_atom)
        return el_refs

    def _create_convhull_data(self):
        """
        Make data suitable for convex hull procedure from the list of entries.
        The procedure is as follows:

        1. First find the elemental references, i.e., the lowest energy entry
           for the vertices of the phase diagram. Using the Li-Fe-O phase
           diagram as an example, this means the lowest energy Li, Fe, and O
           phases.
        2. Calculate the formation energies from these elemental references for
           all entries. Exclude all positive formation energy ones from the
           data for convex hull.
        3. Generate the convex hull data.
        """
        logger.debug("Creating convex hull data...")
        # Remove positive formation energy entries
        qhull_entries = []
        for entry in self.all_entries:
            if self.get_form_energy(entry) <= -self.formation_energy_tol:
                qhull_entries.append(entry)
            else:
                logger.debug("Removing positive formation energy entry " +
                             "{}".format(entry))
        qhull_entries.extend(self.el_refs.values())

        return qhull_entries, self._process_entries_qhulldata(qhull_entries)

    def __repr__(self):
        return self.__str__()

    def __str__(self):
        symbols = [el.symbol for el in self.elements]
        output = ["{} phase diagram".format("-".join(symbols)),
                  "{} stable phases: ".format(len(self.stable_entries)),
                  ", ".join([entry.name
                             for entry in self.stable_entries])]
        return "\n".join(output)

class GrandPotentialPhaseDiagram(PhaseDiagram):
    """
    A class representing a Grand potential phase diagram. Grand potential phase
    diagrams are essentially phase diagrams that are open to one or more
    components. To construct such phase diagrams, the relevant free energy is
    the grand potential, which can be written as the Legendre transform of the
    Gibbs free energy as follows

    Grand potential = G - u\ :sub:`X` N\ :sub:`X`\

    The algorithm is based on the work in the following papers:

    1. S. P. Ong, L. Wang, B. Kang, and G. Ceder, Li-Fe-P-O2 Phase Diagram from
       First Principles Calculations. Chem. Mater., 2008, 20(5), 1798-1807.
       doi:10.1021/cm702327g

    2. S. P. Ong, A. Jain, G. Hautier, B. Kang, G. Ceder, Thermal stabilities
       of delithiated olivine MPO4 (M=Fe, Mn) cathodes investigated using first
       principles calculations. Electrochem. Comm., 2010, 12(3), 427-430.
       doi:10.1016/j.elecom.2010.01.010
    """

    def __init__(self, entries, chempots, elements=None):
        """
        Standard constructor for grand potential phase diagram.

        Args:
            entries:
                A list of PDEntry-like objects having an energy,
                energy_per_atom and composition.
            chempots:
                A dict of {element: float} to specify the chemical potentials
                of the open elements.
            elements:
                Optional list of elements in the phase diagram. If set to None,
                the elements are determined from the entries themselves.
        """
        if elements is None:
            elements = set()
            map(elements.update, [entry.composition.elements
                                  for entry in entries])

        elements = set(elements).difference(chempots.keys())
        all_entries = [GrandPotPDEntry(e, chempots)
                       for e in entries
                       if (not e.is_element) or
                       e.composition.elements[0] in elements]
        self.chempots = chempots

        super(GrandPotentialPhaseDiagram, self).__init__(all_entries, elements)

    def __str__(self):
        output = []
        chemsys = "-".join([el.symbol for el in self.elements])
        output.append("{} grand potential phase diagram with ".format(chemsys))
        output[-1] += ", ".join(["u{}={}".format(el, v)
                                 for el, v in self.chempots.items()])
        output.append("{} stable phases: ".format(len(self.stable_entries)))
        output.append(", ".join([entry.name
                                 for entry in self.stable_entries]))
        return "\n".join(output)

class CompoundPhaseDiagram(PhaseDiagram):
    """
    Generates phase diagrams from compounds as terminations instead of
    elements.
    """

    # Tolerance for determining if amount of a composition is positive.
    amount_tol = 1e-5

    def __init__(self, entries, terminal_compositions,
                 normalize_terminal_compositions=True):
        """
        Args:
            entries:
                Sequence of input entries. For example, if you want a Li2O-P2O5
                phase diagram, you might have all Li-P-O entries as an input.
            terminal_compositions:
                Terminal compositions of phase space. In the Li2O-P2O5 example,
                these will be the Li2O and P2O5 compositions.
            normalize_terminal_compositions:
                Whether to normalize the terminal compositions to a per atom
                basis. If normalized, the energy above hulls will be consistent
                for comparison across systems. Non-normalized terminals are
                more intuitive in terms of compositional breakdowns.
        """
        self.original_entries = entries
        self.terminal_compositions = terminal_compositions
        self.normalize_terminals = normalize_terminal_compositions
        (pentries, species_mapping) = \
            self.transform_entries(entries, terminal_compositions)
        self.species_mapping = species_mapping
        PhaseDiagram.__init__(self, pentries,
                              elements=species_mapping.values())

    def transform_entries(self, entries, terminal_compositions):
        """
        Method to transform all entries to the composition coordinate in the
        terminal compositions. If the entry does not fall within the space
        defined by the terminal compositions, they are excluded. For example,
        Li3PO4 is mapped into a Li2O:1.5, P2O5:0.5 composition. The terminal
        compositions are represented by DummySpecies.

        Args:
            entries:
                Sequence of all input entries
            terminal_compositions:
                Terminal compositions of phase space.

        Returns:
            Sequence of TransformedPDEntries falling within the phase space.
        """
        new_entries = []
        if self.normalize_terminals:
            fractional_comp = [c.get_fractional_composition()
                               for c in terminal_compositions]
        else:
            fractional_comp = terminal_compositions

        #Map terminal compositions to unique dummy species.
        sp_mapping = collections.OrderedDict()
        for i, comp in enumerate(fractional_comp):
            sp_mapping[comp] = DummySpecie("X" + chr(102 + i))

        for entry in entries:
            try:
                rxn = Reaction(fractional_comp, [entry.composition])
                rxn.normalize_to(entry.composition)
                #We only allow reactions that have positive amounts of
                #reactants.
                if all([rxn.get_coeff(comp) <= CompoundPhaseDiagram.amount_tol
                        for comp in fractional_comp]):
                    newcomp = {sp_mapping[comp]: -rxn.get_coeff(comp)
                               for comp in fractional_comp}
                    newcomp = {k: v for k, v in newcomp.items()
                               if v > CompoundPhaseDiagram.amount_tol}
                    transformed_entry = \
                        TransformedPDEntry(Composition(newcomp), entry)
                    new_entries.append(transformed_entry)
            except ReactionError:
                #If the reaction can't be balanced, the entry does not fall
                #into the phase space. We ignore them.
                pass
        return new_entries, sp_mapping

class PhaseDiagramError(Exception):
    """
    An exception class for Phase Diagram.
    """
    def __init__(self, msg):
        """
        Args:
            msg:
                The error message.
        """
        self.msg = msg

    def __str__(self):
        return self.msg


#!/usr/bin/env python
"""
This module contains functions and class definitions for running forward
models of models based on logistic regression.
"""

# stdlib imports
import numpy as np
import os.path
import re
import collections
import shutil
import tempfile
from timeit import default_timer as timer

# third party imports
from mapio.shake import ShakeGrid
from mapio.shake import getHeaderData
from mapio.gmt import GMTGrid
from mapio.gdal import GDALGrid
from mapio.grid2d import Grid2D
from mapio.geodict import GeoDict

from gfail.temphdf import TempHdf
from gfail.spatial import quickcut, trim_ocean
from gfail.utilities import getFileType
from gfail.stats import get_rangebeta

# temporary until mapio is updated
import warnings
warnings.filterwarnings('ignore')

PARAM_PATTERN = 'b[0-9]+'
LAYER_PATTERN = '_layer'
TERM_PATTERN = 'term'

SM_TERMS = ['MW', 'YEAR', 'MONTH', 'DAY', 'HOUR', 'pga', 'pgv', 'mmi']
SM_GRID_TERMS = ['pga', 'pgv', 'mmi']
# these will get np. prepended
OPERATORS = ['log', 'log10', 'arctan', 'power', 'sqrt', 'minimum', 'pi']
FLOATPAT = r'[+-]?(?=\d*[.eE])(?=\.?\d)\d*\.?\d*(?:[eE][+-]?\d+)?'
INTPAT = '[0-9]+'
OPERATORPAT = r'[\+\-\*\/]*'
MONTHS = ['Jan', 'Feb', 'Mar', 'Apr', 'May', 'Jun', 'Jul', 'Aug', 'Sep', 'Oct',
          'Nov', 'Dec']

class LogisticModel(object):
    def __init__(self, shakefile, config, uncertfile=None, saveinputs=False,
                 slopefile=None, bounds=None, slopemod=None,
                 trimfile=None):
        """
        Sets up the logistic model

        Args:
            shakefile (str): Path to shakemap grid.xml file for the event.
            config: configobj object defining the model and its inputs. Only
                one model should be described in each config file.
            uncertfile (str): Path to uncertainty.xml file.
            saveinputs (bool): Save input layers as Grid2D objects in addition
                to the model? If false (the default), it will just output the
                model.
            slopefile (str): Optional path to slopefile that will be resampled
                to the other input files for applying thresholds. OVERWRITES
                VALUE IN CONFIG.
            bounds (dict): Default of None uses ShakeMap boundaries, otherwise
                a dictionary of boundaries to cut to like

                .. code-block:: python

                    bounds = {
                        'xmin': lonmin, 'xmax': lonmax,
                        'ymin': latmin, 'ymax': latmax
                    }
            slopemod (str): How slope input should be modified to be in
                degrees: e.g., ``np.arctan(slope) * 180. / np.pi`` or
                ``slope/100.`` (note that this may be in the config file
                already).
            trimfile (str): shapefile of earth's landmasses to use to cut
                offshore areas.
        """
        mnames = getLogisticModelNames(config)
        if len(mnames) == 0:
            raise Exception('No config file found or problem with config '
                            'file format')
        if len(mnames) > 1:
            raise Exception('Config file contains more than one model which '
                            'is no longer allowed, update your config file '
                            'to the newer format')

        self.model = mnames[0]
        self.config = config
        cmodel = config[self.model]
        self.modeltype = cmodel['gfetype']
        self.coeffs = validateCoefficients(cmodel)
        # key = layer name, value = file name
        self.layers = validateLayers(cmodel)
        self.terms, timeField = validateTerms(cmodel, self.coeffs, self.layers)
        self.interpolations = validateInterpolations(cmodel, self.layers)
        self.units = validateUnits(cmodel)
        self.gmused = [value for term, value in cmodel['terms'].items()
                       if 'pga' in value.lower() or 'pgv' in
                       value.lower() or 'mmi' in value.lower()]
        self.modelrefs, self.longrefs, self.shortrefs = validateRefs(cmodel)
        # self.numstd = numstd
        self.clips = validateClips(cmodel, self.layers, self.gmused)
        self.notes = ''

        if cmodel['baselayer'] not in list(self.layers.keys()):
            raise Exception('You must specify a base layer corresponding to '
                            'one of the files in the layer section.')
        self.saveinputs = saveinputs
        if slopefile is None:
            try:
                self.slopefile = cmodel['slopefile']
            except BaseException:
                # print('Slopefile not specified in config, no slope '
                #      'thresholds will be applied\n')
                self.slopefile = None
        else:
            self.slopefile = slopefile
        if slopemod is None:
            try:
                self.slopemod = cmodel['slopemod']
            except BaseException:
                self.slopemod = None

        # See if trimfile exists
        if trimfile is not None:
            if not os.path.exists(trimfile):
                print(
                    'trimfile defined does not exist: %s\nOcean will not be '
                    'trimmed' % trimfile)
                self.trimfile = None
            elif os.path.splitext(trimfile)[1] != '.shp':
                print('trimfile must be a shapefile, ocean will not be '
                      'trimmed')
                self.trimfile = None
            else:
                self.trimfile = trimfile
        else:
            self.trimfile = None

        # Get month of event
        griddict, eventdict, specdict, fields, uncertainties = \
            getHeaderData(shakefile)
        MONTH = MONTHS[eventdict['event_timestamp'].month - 1]

        # Figure out how/if need to cut anything
        geodict = ShakeGrid.getFileGeoDict(shakefile, adjust='res')
        if bounds is not None:  # Make sure bounds are within ShakeMap Grid
            if geodict.xmin < geodict.xmax:  # only if signs are not opposite
                if (geodict.xmin > bounds['xmin'] or
                        geodict.xmax < bounds['xmax'] or
                        geodict.ymin > bounds['ymin'] or
                        geodict.ymax < bounds['ymax']):
                    print('Specified bounds are outside shakemap area, using '
                          'ShakeMap bounds instead.')
                    bounds = None

        if bounds is not None:
            tempgdict = GeoDict.createDictFromBox(
                bounds['xmin'], bounds['xmax'],
                bounds['ymin'], bounds['ymax'],
                geodict.dx, geodict.dy, inside=False)
            # If Shakemap geodict crosses 180/-180 line, fix geodict so
            # things don't break
            if geodict.xmin > geodict.xmax:
                if tempgdict.xmin < 0:
                    geodict._xmin -= 360.
                else:
                    geodict._xmax += 360.
            gdict = geodict.getBoundsWithin(tempgdict)
        else:
            gdict = geodict

        # Now find the layer that is our base layer and get the largest bounds
        # we can guarantee not to exceed shakemap bounds
        basefile = self.layers[cmodel['baselayer']]
        ftype = getFileType(basefile)
        if ftype == 'esri':
            basegeodict, firstcol = GDALGrid.getFileGeoDict(basefile)
            if basegeodict == gdict:
                sampledict = gdict
            else:
                sampledict = basegeodict.getBoundsWithin(gdict)
        elif ftype == 'gmt':
            basegeodict, firstcol = GMTGrid.getFileGeoDict(basefile)
            if basegeodict == gdict:
                sampledict = gdict
            else:
                sampledict = basegeodict.getBoundsWithin(gdict)
        else:
            raise Exception('All predictor variable grids must be a valid '
                            'GMT or ESRI file type.')

        # Do we need to subdivide baselayer?
        if 'divfactor' in self.config[self.model].keys():
            divfactor = float(self.config[self.model]['divfactor'])
            if divfactor != 1.:
                # adjust sampledict so everything will be resampled
                newxmin = sampledict.xmin - sampledict.dx / \
                    2. + sampledict.dx / (2. * divfactor)
                newymin = sampledict.ymin - sampledict.dy / \
                    2. + sampledict.dy / (2. * divfactor)
                newxmax = sampledict.xmax + sampledict.dx / \
                    2. - sampledict.dx / (2. * divfactor)
                newymax = sampledict.ymax + sampledict.dy / \
                    2. - sampledict.dy / (2. * divfactor)
                newdx = sampledict.dx / divfactor
                newdy = sampledict.dy / divfactor
                if np.abs(newxmax) > 180.:
                    newxmax = np.sign(newxmax) * 180.
                if np.abs(newxmin) > 180.:
                    newxmin = np.sign(newxmin) * 180.

                sampledict = GeoDict.createDictFromBox(
                    newxmin, newxmax, newymin,
                    newymax, newdx, newdy, inside=True)

        # Find slope thresholds, if applicable
        self.slopemin = 'none'
        self.slopemax = 'none'
        if self.slopefile is not None:
            try:
                self.slopemin = float(config[self.model]['slopemin'])
                self.slopemax = float(config[self.model]['slopemax'])
            except BaseException:
                print('Could not find slopemin and/or slopemax in config, '
                      'limits. No slope thresholds will be applied.')
                self.slopemin = 'none'
                self.slopemax = 'none'

        # Make temporary directory for hdf5 pytables file storage
        self.tempdir = tempfile.mkdtemp()

        # now load the shakemap, resampling and padding if necessary
        temp = ShakeGrid.load(shakefile)  # , adjust='res')
        self.shakedict = temp.getShakeDict()
        self.eventdict = temp.getEventDict()
        self.shakemap = {}

        # Read both PGA and PGV in, may need them for thresholds
        for gm in ['pga', 'pgv']:
            junkfile = os.path.join(self.tempdir, 'temp.bil')
            GDALGrid.copyFromGrid(temp.getLayer(gm)).save(junkfile)
            if gm in self.interpolations.keys():
                intermeth = self.interpolations[gm]
            else:
                intermeth = 'bilinear'
            junkgrid = quickcut(junkfile, sampledict, precise=True,
                                method=intermeth, override=True)
            if gm in self.clips:
                junkgrid.setData(np.clip(junkgrid.getData(),
                                         self.clips[gm][0], self.clips[gm][1]))
            self.shakemap[gm] = TempHdf(
                junkgrid, os.path.join(self.tempdir, '%s.hdf5' % gm))
            os.remove(junkfile)
        del(temp)

        # get updated geodict
        sampledict = junkgrid.getGeoDict()

        # take uncertainties into account, if available
        if uncertfile is not None:
            self.uncert = {}
            # try:
            # Only read in the ones that will be needed
            temp = ShakeGrid.load(uncertfile)
            already = []
            for gm in self.gmused:
                if 'pgv' in gm:
                    gmsimp = 'pgv'
                elif 'pga' in gm:
                    gmsimp = 'pga'
                elif 'mmi' in gm:
                    gmsimp = 'mmi'
                if gmsimp in already:
                    continue
                junkfile = os.path.join(self.tempdir, 'temp.bil')
                GDALGrid.copyFromGrid(temp.getLayer(
                    'std%s' % gmsimp)).save(junkfile)
                if gmsimp in self.interpolations.keys():
                    intermeth = self.interpolations[gmsimp]
                else:
                    intermeth = 'bilinear'
                junkgrid = quickcut(junkfile, sampledict, precise=True,
                                    method=intermeth, override=True)
                if gmsimp in self.clips:
                    junkgrid.setData(
                        np.clip(junkgrid.getData(), self.clips[gmsimp][0],
                                self.clips[gmsimp][1]))
                self.uncert['std' + gmsimp] = TempHdf(
                    junkgrid, os.path.join(self.tempdir,
                                           'std%s.hdf5' % gmsimp))
                already.append(gmsimp)
                os.remove(junkfile)
            del(temp)
            # except:
            # print('Could not read uncertainty file, ignoring '
            #       'uncertainties')
            # self.uncert = None
        else:
            self.uncert = None

        # Load the predictor layers, save as hdf5 temporary files, put file
        # locations into a dictionary.

        # Will be replaced in the next section if a slopefile was defined
        self.nonzero = None

        # key = layer name, value = grid object
        self.layerdict = {}

        didslope = False
        for layername, layerfile in self.layers.items():
            start = timer()
            if isinstance(layerfile, list):
                for lfile in layerfile:
                    if timeField == 'MONTH':
                        if lfile.find(MONTH) > -1:
                            layerfile = lfile
                            # ftype = getFileType(layerfile)
                            interp = self.interpolations[layername]
                            temp = quickcut(layerfile, sampledict,
                                            precise=True, method=interp)
                            if layername in self.clips:
                                temp.setData(
                                    np.clip(temp.getData(),
                                            self.clips[layername][0],
                                            self.clips[layername][1]))
                            self.layerdict[layername] = TempHdf(
                                temp, os.path.join(self.tempdir,
                                                   '%s.hdf5' % layername))
                            del(temp)
            else:
                interp = self.interpolations[layername]
                temp = quickcut(layerfile, sampledict,
                                precise=True, method=interp)
                if layername in self.clips:
                    temp.setData(
                        np.clip(temp.getData(),
                                self.clips[layername][0],
                                self.clips[layername][1]))
                # Convert unconsolidated sediments to more reasonable coeff
                if layername == 'rock':
                    sub1 = temp.getData()
                    # Change to mixed sed rock coeff
                    sub1[sub1 <= -3.21] = -1.36
                    temp.setData(sub1)
                    self.notes += 'unconsolidated sediment coefficient ' \
                                  'changed to -1.36 (weaker) from -3.22 to ' \
                                  'better reflect that this ' \
                                  'unit is not actually strong\n'
                self.layerdict[layername] = TempHdf(
                    temp, os.path.join(self.tempdir, '%s.hdf5' % layername))
                td = temp.getGeoDict()
                if td != sampledict:
                    raise Exception(
                        'Geodictionaries of resampled files do not match')

                if layerfile == self.slopefile:
                    flag = 0
                    if self.slopemin == 'none' and self.slopemax == 'none':
                        flag = 1
                    if self.slopemod is None:
                        slope1 = temp.getData().astype(float)
                        slope = 0
                    else:
                        try:
                            slope = temp.getData().astype(float)
                            slope1 = eval(self.slopemod)
                        except BaseException:
                            print('slopemod provided not valid, continuing '
                                  'without slope thresholds.')
                            flag = 1
                    if flag == 0:
                        nonzero = np.array(
                            [(slope1 > self.slopemin) &
                             (slope1 <= self.slopemax)])
                        self.nonzero = nonzero[0, :, :]
                        del(slope1)
                        del(slope)
                    else:
                        # Still remove areas where the slope equals exactly
                        # 0.0 to remove offshore liq areas.
                        nonzero = np.array([slope1 != 0.0])
                        self.nonzero = nonzero[0, :, :]
                        del(slope1)
                    didslope = True
                del(temp)

            print('Loading %s layer: %1.1f sec'
                  % (layername, timer() - start))

        if didslope is False and self.slopefile is not None:
            # Slope didn't get read in yet
            temp = quickcut(self.slopefile, sampledict, precise=True,
                            method='bilinear')
            flag = 0
            if self.slopemin == 'none' and self.slopemax == 'none':
                flag = 1
            if self.slopemod is None:
                slope1 = temp.getData().astype(float)
                slope = 0
            else:
                try:
                    slope = temp.getData().astype(float)
                    slope1 = eval(self.slopemod)
                except BaseException:
                    print('slopemod provided not valid, continuing without '
                          'slope thresholds')
                    flag = 1
            if flag == 0:
                nonzero = np.array([(slope1 > self.slopemin) &
                                    (slope1 <= self.slopemax)])
                self.nonzero = nonzero[0, :, :]
                del(slope1)
                del(slope)
            else:
                # Still remove areas where the slope equals exactly
                # 0.0 to remove offshore liq areas.
                nonzero = np.array([slope1 != 0.0])
                self.nonzero = nonzero[0, :, :]
                del(slope1)

        self.nuggets = [str(self.coeffs['b0'])]

        ckeys = sorted(self.terms.keys())
        for key in ckeys:
            term = self.terms[key]
            coeff = self.coeffs[key]
            self.nuggets.append('(%g * %s)' % (coeff, term))

        self.equation = ' + '.join(self.nuggets)
        self.geodict = sampledict

    def getEquations(self):
        """
        Method for LogisticModel class to extract strings defining the
        equations for the model for median ground motions.

        Returns:
            equation: the equation for median ground motions,

        """
        return self.equation

    def getGeoDict(self):
        """
        Returns the geodictionary of the LogisticModel class defining bounds
        and resolution of model inputs and outputs.

        Returns:
            geodict: mapio geodict object
        """
        return self.geodict

    def calculate(self, cleanup=True, rowmax=300, colmax=None):
        """
        Calculate the model.

        Args:
            cleanup (bool): If True, delete temporary hdf5 files
            rowmax (int): Number of rows to compute at once; If None, all rows
                will be computed at once.
            colmax (int): Number of columns to compute at once; If None, all
                columns will be computed at once.
        Returns:
            dict: Dictionary containing the model results (and model inputs if
            saveinputs was set to True). See
            `the description <https://github.com/usgs/groundfailure#api-for-model-output>`_
            of the structure.
        """
        tk = list(self.shakemap.keys())[0]
        # Figure out what slices to do
        rowstarts, rowends, colstarts, colends = \
            self.shakemap[tk].getSliceDiv(rowmax, colmax)

        # Make empty matrix to fill
        X = np.empty([self.geodict.ny, self.geodict.nx])

        # Loop through slices, appending output each time
        for rowstart, rowend, colstart, colend in \
                zip(rowstarts, rowends, colstarts, colends):
            X[rowstart:rowend, colstart:colend] = eval(self.equation)

        P = 1 / (1 + np.exp(-X))

        if 'vs30max' in self.config[self.model].keys():
            vs30 = self.layerdict['vs30'].getSlice(
                None, None, None, None, name='vs30')
            P[vs30 > float(self.config[self.model]['vs30max'])] = 0.0

        if 'minpgv' in self.config[self.model].keys():
            pgv = self.shakemap['pgv'].getSlice(
                None, None, None, None, name='pgv')
            P[pgv < float(self.config[self.model]['minpgv'])] = 0.0

        if 'minpga' in self.config[self.model].keys():
            pga = self.shakemap['pga'].getSlice(
                None, None, None, None, name='pga')
            P[pga < float(self.config[self.model]['minpga'])] = 0.0

        if self.uncert is not None:  # hard code for now
            if 'Zhu and others (2017)' in self.modelrefs['shortref']:
                if 'stddev' in self.layerdict.keys():
                    stdX = self.layerdict['stddev'].getSlice()
                else:
                    stdX = float(self.config[self.model]['default_stddev'])
                varX = stdX**2. + \
                    (self.coeffs['b1']**2. *
                     self.uncert['stdpgv'].getSlice()**2.)
                varP = (np.exp(-X) / (np.exp(-X) + 1)**2.)**2. * varX
                if 'coverage' in self.config[self.model].keys():
                    a = 0.4915
                    b = 42.4
                    c = 9.165
                    # ((2*a*b*c*np.exp(2*c*P))/(b+np.exp(c*P))**3.)**2.*varP
                    varL = ((2 * a * b * c * np.exp(-c * P)) /
                            ((1 + b * np.exp(-c * P))**3.))**2. * varP
                    std1 = np.sqrt(varL)
                else:
                    std1 = np.sqrt(varP)
            elif 'Jessee' in self.modelrefs['shortref']:
                if 'stddev' in self.layerdict.keys():
                    stdX = self.layerdict['stddev'].getSlice()
                else:
                    stdX = float(self.config[self.model]['default_stddev'])
                cfs = self.coeffs
                slp = self.layerdict['slope']
                std = self.uncert['stdpgv']
                varX = stdX**2. + ((
                    cfs['b1'] + cfs['b6'] *
                    (np.arctan(slp.getSlice()) * 180 / np.pi))**2.
                    * std.getSlice()**2.)
                varP = (np.exp(-X) / (np.exp(-X) + 1)**2.)**2. * varX
                if 'coverage' in self.config[self.model].keys():
                    a = -7.592
                    b = 5.237
                    c = -3.042
                    d = 4.035
                    varL = (np.exp(a + b * P + c * P**2. + d * P**3.) *
                            (b + 2. * P * c + 3. * d * P**2.))**2. * varP
                    std1 = np.sqrt(varL)
                else:
                    std1 = np.sqrt(varP)
            else:
                print('cannot do uncertainty for %s model, skipping' %
                      self.modelrefs['shortref'])
                self.uncert = None
                std1 = None
        else:
            std1 = None

        # P needs to be converted to areal coverage AFTER dealing with uncert
        if 'coverage' in self.config[self.model].keys():
            eqn = self.config[self.model]['coverage']['eqn']
            P = eval(eqn)

        # Compute quantiles
        compute_quantiles = False
        mconf = self.config[self.model]
        if(('conf_int_probabilities' in mconf) and
                (std1 is not None)):
            compute_quantiles = True
            ci_probabilities = [
                float(cip) for cip in mconf['conf_int_probabilities']]

        if compute_quantiles:
            quantile_dict = {}
            pmax = float(mconf['maxprob'])
            beta_p = P / pmax * (((pmax * P - P**2) / std1**2) - 1)
            beta_q = (1 - P / pmax) * (((pmax * P - P**2) / std1**2) - 1)
            for ci_prob in ci_probabilities:
                min_quantile = str(np.round(100 * (1.0 - ci_prob) / 2.0, 1))
                max_quantile = str(np.round(
                    100 * (1 - ((1.0 - ci_prob)) / 2.0), 1))
                min_prob, max_prob = get_rangebeta(
                    beta_p, beta_q, ci_prob, minlim=0, maxlim=pmax)
                quantile_dict[min_quantile] = min_prob
                quantile_dict[max_quantile] = max_prob

        if self.slopefile is not None and self.nonzero is not None:
            # Apply slope min/max limits
            print('applying slope thresholds')
            P = P * self.nonzero
            if std1 is not None:
                # No uncert for masked values
                std1[P == 0] = 0.
                if compute_quantiles:
                    for q in quantile_dict.values():
                        q[P == 0] = 0.

        # Stuff into Grid2D object
        if 'Jessee' in self.modelrefs['shortref']:
            if 'coverage' not in self.config[self.model].keys():
                units5 = 'Relative Hazard'
            else:
                units5 = 'Proportion of area affected'
        elif 'Zhu' in self.modelrefs['shortref']:
            if 'coverage' not in self.config[self.model].keys() and \
                    '2017' in self.modelrefs['shortref']:
                units5 = 'Relative Hazard'
            else:
                units5 = 'Proportion of area affected'
        else:
            units5 = 'Probability of any occurrence'

        shakedetail = (
            '%s_ver%s'
            % (self.shakedict['shakemap_id'],
               self.shakedict['shakemap_version']))
        description = {
            'name': self.modelrefs['shortref'],
            'longref': self.modelrefs['longref'],
            'units': units5,
            'shakemap': shakedetail,
            'event_id': self.eventdict['event_id'],
            'parameters': {'slopemin': self.slopemin,
                           'slopemax': self.slopemax,
                           'modeltype': self.modeltype,
                           'notes': self.notes}}
        if 'vs30max' in self.config[self.model].keys():
            description['vs30max'] = float(self.config[self.model]['vs30max'])
        if 'minpgv' in self.config[self.model].keys():
            description['minpgv'] = float(self.config[self.model]['minpgv'])

        Pgrid = Grid2D(P, self.geodict)
        if self.trimfile is not None:
            # Turn all offshore cells to nan
            Pgrid = trim_ocean(Pgrid, self.trimfile)
        rdict = collections.OrderedDict()
        rdict['model'] = {
            'grid': Pgrid,
            'label': '%s estimate - %s' % (self.modeltype.capitalize(),
                                           units5.title()),
            'type': 'output',
            'description': description
        }
        if self.uncert is not None:
            Stdgrid = Grid2D(std1, self.geodict)
            if self.trimfile is not None:
                Stdgrid = trim_ocean(
                    Stdgrid, self.trimfile)
            rdict['std'] = {
                'grid': Stdgrid,
                'label': ('%s estimate - %s (std)'
                          % (self.modeltype.capitalize(),
                             units5.title())),
                'type': 'output',
                'description': description
            }
            if compute_quantiles:
                for quantile, qgrid in quantile_dict.items():
                    Qgrid = Grid2D(qgrid, self.geodict)
                    qname = "quantile%s" % quantile
                    rdict[qname] = {
                        'grid': Qgrid,
                        'label': (
                            '%s %sth percentile - %s'
                            % (self.modeltype.capitalize(), quantile,
                               units5.title())),
                        'type': 'output',
                        'description': description
                    }

        # This step might swamp memory for higher resolution runs
        if self.saveinputs is True:
            for layername, layergrid in list(self.layerdict.items()):
                units = self.units[layername]
                if units is None:
                    units = ''
                rdict[layername] = {
                    'grid': Grid2D(
                        layergrid.getSlice(
                            None, None, None, None, name=layername),
                        self.geodict
                    ),
                    'label': '%s (%s)' % (layername, units),
                    'type': 'input',
                    'description': {
                        'units': units,
                        'name': self.shortrefs[layername],
                        'longref': self.longrefs[layername]
                    }
                }
            for gmused in self.gmused:
                if 'pga' in gmused:
                    units = '%g'
                    getkey = 'pga'
                elif 'pgv' in gmused:
                    units = 'cm/s'
                    getkey = 'pgv'
                elif 'mmi' in gmused:
                    units = 'intensity'
                    getkey = 'mmi'
                else:
                    continue
                    # Layer is derived from several input layers, skip
                    # outputting this layer

                if getkey in rdict:
                    continue

                layer = self.shakemap[getkey].getSlice(
                    None, None, None, None, name=getkey)
                rdict[getkey] = {
                    'grid': Grid2D(layer, self.geodict),
                    'label': '%s (%s)' % (getkey.upper(), units),
                    'type': 'input',
                    'description': {
                        'units': units,
                        'shakemap': shakedetail
                    }
                }
        if cleanup:
            shutil.rmtree(self.tempdir)
        return rdict

def getLogisticModelNames(config):
    """
    Get the names of the models present in the configobj

    Args:
        config: configobj object defining the model and its inputs.

    Returns:
        list: list of model names.
    """
    names = []
    lmodel_space = config
    for key, value in lmodel_space.items():
        if isinstance(value, str):
            continue
        else:  # this is a model
            names.append(key)
    return names

def getAllGridFiles(indir):
    """
    Get list of all gmt or esri (.grd, .bil) files in a directory.

    Args:
        indir (str): Directory to search.
    Returns:
        list: List of file names.
    """
    # TODO MOVE TO MAPIO
    tflist = os.listdir(indir)
    flist = []
    for tf in tflist:
        fullfile = os.path.join(indir, tf)
        ftype = getFileType(fullfile)
        if ftype in ['gmt', 'esri']:
            flist.append(fullfile)
    return flist

def validateCoefficients(cmodel):
    """
    Ensures coefficients provided in model description are valid and outputs
    a dictionary of the coefficients.

    Args:
        cmodel (dict): Sub-dictionary from config for specific model,
            for example:

            .. code-block:: python

                cmodel = config['test_model']

    Returns:
        dict: a dictionary of model coefficients named b0, b1, b2...
    """
    coeffs = {}
    for key, value in cmodel['coefficients'].items():
        if re.search('b[0-9]*', key) is None:
            raise Exception('coefficients must be named b0, b1, ...')
        coeffs[key] = float(value)
    if 'b0' not in list(coeffs.keys()):
        raise Exception('coefficients must include an intercept '
                        'coefficient named b0.')
    return coeffs

def validateClips(cmodel, layers, gmused):
    """
    Ensures coefficients provided in model description are valid and outputs
    a dictionary of the coefficients.

    Args:
        cmodel (dict): Sub-dictionary from config for specific model,
            for example:

            .. code-block:: python

                cmodel = config['test_model']
        layers: dictionary of layer names
        gmused (list): List of ground motion parameters used

    Returns:
        dict: a dictionary of clip values for each layer (if exists)
    """
    clips = {}
    if 'clip' in cmodel:
        for key, value in cmodel['clip'].items():
            if key not in layers:
                if key not in gmused:
                    x1 = [par for par in gmused if key in par]
                    if len(x1) == 0:
                        raise Exception(
                            'Clipping key %s does not match any layers'
                            % key)
            clips[key] = (float(value[0]), float(value[1]))
    return clips

def validateLayers(cmodel):
    """
    Ensures all input files required to run the model exist and are valid
    file types. Make sure all layers are available for area of run

    Args:
        cmodel (dict): Sub-dictionary from config for specific model,
            for example,

            .. code-block:: python

                cmodel = config['test_model']

    Returns:
        dict: a dictionary of file names, e.g.

        .. code-block:: python

            {
                'slope': 'slopefile.bil',
                'vs30': 'vs30.grd'
            }

    """
    layers = {}
    longrefs = {}
    shortrefs = {}
    for key in cmodel['layers'].keys():
        for item, value in cmodel['layers'][key].items():
            if item == 'file':
                ftype = getFileType(value)
                if ftype == 'unknown':
                    raise Exception('layer file %s is not a valid GMT or '
                                    'ESRI file.' % value)
                if ftype == 'dir':
                    value = getAllGridFiles(value)
                layers[key] = value
            elif item == 'shortref':
                shortrefs[key] = value
            elif item == 'longref':
                longrefs[key] = value
    return layers

def validateTerms(cmodel, coeffs, layers):
    """
    Reformats model inputs from config file, replacing functions with numpy
    functions, inserting code for extracting data from each layer (required
    to run eval in the calculate step), addressing any time variables, and
    checks that term names match coefficient names.

    Args:
        cmodel (dict): Sub-dictionary from config for specific model,
            e.g.

            .. code-block:: python

                cmodel = config['test_model']

        coeffs (dict): Dictionary of model coefficients, e.g.

            .. code-block:: python

                {'b0': 3.5, 'b1': -0.01}

        layers (dict): Dictionary of file names for all input layers, e.g.

            .. code-block:: python

                {'slope': 'slopefile.bil', 'vs30': 'vs30.grd'}

    Returns:
        tuple: (terms, timeField), where
            - 'terms' is a dictionary of terms that form the model equation,
              e.g.

            .. code-block:: python

                {
                    'b1': "self.layerdict['friction'].getData()",
                    'b2': "self.layerdict['slope'].getData()/100."
                }

            - 'timeField' indicates the time that is used to know which input
              file to read in, e.g. for monthly average precipitation, 'MONTH'.
    """
    # TODO:
    #    - Return a time field for every term, not just one global one.

    terms = {}
    timeField = None
    for key, value in cmodel['terms'].items():
        if key not in list(coeffs.keys()):
            raise Exception('Term names must match names of coefficients')
        # replace log with np.log, make sure variables are all in layers list,
        # etc.
        term, rem, tTimeField = checkTerm(value, layers)
        if tTimeField is not None:
            timeField = tTimeField
        if len(rem):
            msg = ('Term "%s" contains the unknown text fragment "%s". '
                   'This may cause the expression to fail.')
            tpl = (term, rem)
            raise Exception(msg % tpl)
        terms[key] = term
    return terms, timeField

def validateInterpolations(cmodel, layers):
    """Validate logistic model interpolation.

    Args:
        cmodel (dict): Sub-dictionary from config for specific model.
        layers (dict): Dictionary of file names for all input layers.

    Returns:
        dict: Model interpolation methods.
    """
    interpolations = {}
    for key, value in cmodel['interpolations'].items():
        if key not in list(layers.keys()):
            raise Exception(
                'Interpolation key %s does not match any names of layers'
                % key)
        methods = ['linear', 'nearest', 'cubic', 'bilinear']
        if value not in methods:
            raise Exception(
                'Interpolation method %s not in approved list of methods: %s'
                % (key, str(methods)))
        interpolations[key] = value
    for key in list(layers.keys()):
        if key not in list(interpolations.keys()):
            raise Exception(
                'No interpolation method configured for layer %s' % key)
    return interpolations

def validateUnits(cmodel):
    """Validate model units.

    Args:
        cmodel (dict): Sub-dictionary from config for specific model.

    Returns:
        dict: Model units.
    """
    units = {}
    for key in cmodel['layers'].keys():
        if 'units' in cmodel['layers'][key]:
            units[key] = cmodel['layers'][key]['units']
        else:
            raise Exception('No unit string configured for layer %s' % key)
    return units

def validateRefs(cmodel):
    """Validate references for models and layers.

    Args:
        cmodel (dict): Sub-dictionary from config for specific model.

    Returns:
        tuple: (modelrefs, longrefs, shortrefs) where:
            * modelrefs: dictionary of citation information for model
                keys='longref', 'shortref'
            * shortrefs: dictionary containing short reference for each
                input layer
            * longrefs: dictionary containing full references for each
                input layer

    """
    longrefs = {}
    shortrefs = {}
    modelrefs = {}
    for key in cmodel['layers'].keys():
        if 'longref' in cmodel['layers'][key]:
            longrefs[key] = cmodel['layers'][key]['longref']
        else:
            print('No longref provided for layer %s' % key)
            longrefs[key] = 'unknown'
        if 'shortref' in cmodel['layers'][key]:
            shortrefs[key] = cmodel['layers'][key]['shortref']
        else:
            print('No shortref provided for layer %s' % key)
            shortrefs[key] = 'unknown'
    try:
        modelrefs['longref'] = cmodel['longref']
    except BaseException:
        print('No model longref provided')
        modelrefs['longref'] = 'unknown'
    try:
        modelrefs['shortref'] = cmodel['shortref']
    except BaseException:
        print('No model shortref provided')
        modelrefs['shortref'] = 'unknown'
    return modelrefs, longrefs, shortrefs

def checkTerm(term, layers):
    """Checks terms of equation and replaces text with machine readable
    operators

    Args:
        term: term from model configuration file
        layers: dictionary of file names for all input layers

    Returns:
        tuple: (term, tterm, timeField) where:
            * term: dictionary of verified terms for equation with keys
                corresponding to each layer name
            * tterm: any unconverted and unverified text that may cause
                expression to fail
            * timeField: if any inputs are time dependent, output is unit of
                time (e.g., 'YEAR'), otherwise, None.
    """
    # startterm = term
    # Strip out everything that isn't: 0-9.() operators, +-/* or layer names.
    # Anything left is an unknown symbol.
    tterm = term
    # remove log, sqrt, etc.
    for op in OPERATORS:
        tterm = tterm.replace(op, '')
    # remove ShakeMap variables
    for sm_term in SM_TERMS:
        tterm = tterm.replace(sm_term, '')
    # remove layer names
    for layer in layers:
        tterm = tterm.replace(layer, '')
    # remove arithmetic operators
    tterm = re.sub(OPERATORPAT, '', tterm)
    # remove floating point numbers
    tterm = re.sub(FLOATPAT, '', tterm)
    # remove integer numbers
    tterm = re.sub(INTPAT, '', tterm)
    # remove parentheses
    tterm = re.sub('[()]*', '', tterm)
    # remove any blank spaces
    tterm = tterm.strip()
    # remove commas
    tterm = tterm.strip(',')
    # anything left *might* cause an error
    for op in OPERATORS:
        if term.find(op) > -1:
            term = term.replace(op, 'np.' + op)

    for sm_term in SM_GRID_TERMS:
        term = term.replace(
            sm_term,
            "self.shakemap['%s'].getSlice(rowstart, rowend, "
            "colstart, colend, name='%s')" % (sm_term, sm_term))

    # replace the macro MW with the magnitude value from the shakemap
    term = term.replace('MW', "self.eventdict['magnitude']")

    # term.replace('YEAR',"self.shakemap.getEventDict()['event_time'].year")
    # hasTime = False
    timeField = None
    for unit in ['YEAR', 'MONTH', 'DAY', 'HOUR']:
        if term.find(unit) > -1:
            term = term.replace(unit, '')
            timeField = unit

    for layer in layers:
        if layer == 'friction':
            term = term.replace(
                layer,
                "np.nan_to_num(self.layerdict['%s'].getSlice(rowstart, "
                "rowend, colstart, colend, name='%s'))" % (layer, layer))
        else:
            term = term.replace(
                layer,
                "self.layerdict['%s'].getSlice(rowstart, rowend, colstart, "
                "colend, name='%s')" % (layer, layer))
    return term, tterm, timeField


import logging

import numpy as np

from renormalizer.mps import MpDm, Mpo
from renormalizer.mps.tdh import unitary_propagation
from renormalizer.utils import TdMpsJob, Quantity, EvolveConfig
from renormalizer.property import Property

logger = logging.getLogger(__name__)

class ThermalProp(TdMpsJob):
    r"""
    Thermally Propagate initial matrix product density operator (:class:`~renormalizer.mps.MpDm`) in imaginary time.

    Args:
        init_mpdm (:class:`~renormalizer.mps.MpDm`): the initial density matrix to be propagated. Usually identity.
        h_mpo (:class:`~renormalizer.mps.Mpo`): the system Hamiltonian.
        exact (bool): whether propagate assuming Hamiltonian is local
            :math:`\hat H = \sum_i \hat H_i = \sum_{in} \omega_{in} b^\dagger_{in} b_{in}` and
            exact propagation is possible through :math:`e^{xH} = e^{xh_1}e^{xh_2} \cdots e^{xh_n}`.
            If set to ``True``, properties such as occupations are not calculated during
            time evolution for better efficiency.
        space (str): the space of exact propagation. Possible options are ``"GS"`` or ``"EX"``.
            If set to ``"GS"``, then the exact propagation is performed in zero exciton space.
            If set to ``"EX"``, then the exact propagation is performed in one exciton space,
            i.e. the vibrations are regarded as displaced oscillators.
        evolve_config (:class:`~renormalizer.utils.EvolveConfig`): config when evolving the MpDm in imaginary time.
        dump_mps (bool): if dump mps when dumping
        dump_dir (str): the directory for logging and numerical result output.
        job_name (str): the name of the calculation job which determines the file name of the logging and numerical result output.
        properties (:class:`~renormalizer.property.Property`) calculate other properties with interface in Property
    """
    def __init__(
        self,
        init_mpdm: MpDm,
        h_mpo: Mpo,
        exact: bool = False,
        space: str = "GS",
        evolve_config: EvolveConfig = None,
        dump_mps: bool = False,
        dump_dir: str = None,
        job_name: str = None,
        properties: Property = None,
        auto_expand: bool = True,
    ):
        self.init_mpdm: MpDm = init_mpdm.canonicalise()
        self.h_mpo = h_mpo
        self.exact = exact
        assert space in ["GS", "EX"]
        self.space = space
        self.energies = []
        self._e_occupations_array = []
        self._ph_occupations_array = []
        self._vn_entropy_array = []
        self.properties = properties
        self.auto_expand = auto_expand

        super().__init__(evolve_config=evolve_config, dump_mps=dump_mps, dump_dir=dump_dir,
                job_name=job_name)

    def init_mps(self):
        self.init_mpdm.evolve_config = self.evolve_config
        if self.evolve_config.is_tdvp and self.auto_expand:
            self.init_mpdm = self.init_mpdm.expand_bond_dimension(self.h_mpo)
        return self.init_mpdm

    def process_mps(self, mps):
        if self.exact:
            # skip the fuss for efficiency
            return
        new_energy = mps.expectation(self.h_mpo)
        self.energies.append(new_energy)
        for attr_str in ["e_occupations", "ph_occupations"]:
            attr = getattr(mps, attr_str)
            logger.info(f"{attr_str}: {attr}")
            self_array = getattr(self, f"_{attr_str}_array")
            self_array.append(attr)
        vn_entropy = mps.calc_vn_entropy()
        self._vn_entropy_array.append(vn_entropy)
        logger.info(f"vn entropy: {vn_entropy}")
        logger.info(
            f"Energy: {new_energy}, total electron: {self._e_occupations_array[-1].sum()}"
        )

        # calculate other properties defined in Property
        if self.properties is not None:
            self.properties.calc_properties(mps)

    def evolve_exact(self, old_mpdm, evolve_dt):
        MPOprop, HAM, Etot = old_mpdm.hybrid_exact_propagator(
            self.h_mpo, evolve_dt.imag, space=self.space
        )
        new_mpdm = MPOprop.apply(old_mpdm)
        unitary_propagation(new_mpdm.tdh_wfns, HAM, Etot, evolve_dt)
        # partition function can't be obtained. It's not practical anyway.
        # The function is too large to be fit into float64 even float128
        new_mpdm.canonicalise(normalize=True)
        new_mpdm.normalize(1.0)
        return new_mpdm

    def evolve_prop(self, old_mpdm, evolve_dt):
        h_mpo = Mpo(self.h_mpo.mol_list, offset=Quantity(self.energies[-1]))
        return old_mpdm.evolve(h_mpo, evolve_dt)

    def evolve_single_step(self, evolve_dt):
        old_mpdm = self.latest_mps
        if self.exact:
            new_mpdm = self.evolve_exact(old_mpdm, evolve_dt)
        else:
            new_mpdm = self.evolve_prop(old_mpdm, evolve_dt)
        return new_mpdm

    def evolve(self, evolve_dt=None, nsteps=None, evolve_time=None):
        if evolve_dt is not None:
            assert np.iscomplex(evolve_dt) and evolve_dt.imag < 0
        if evolve_time is not None:
            assert np.iscomplex(evolve_time) and evolve_time.imag < 0
        super().evolve(evolve_dt, nsteps, evolve_time)

    @property
    def e_occupations_array(self):
        return np.array(self._e_occupations_array)

    @property
    def ph_occupations_array(self):
        return np.array(self._ph_occupations_array)

    @property
    def vn_entropy_array(self):
        return np.array(self._vn_entropy_array)

    def get_dump_dict(self):
        dump_dict = dict()
        dump_dict["time series"] = [-t.imag for t in self.evolve_times]
        dump_dict["energies"] = self.energies
        dump_dict["electron occupations array"] = self.e_occupations_array.tolist()
        dump_dict["phonon occupations array"] = self.ph_occupations_array.tolist()
        dump_dict["vn entropy array"] = self.vn_entropy_array.tolist()

        if self.properties is not None:
            for prop_str in self.properties.prop_res.keys():
                dump_dict[prop_str] = self.properties.prop_res[prop_str]

        return dump_dict

def load_thermal_state(mol_list, path: str):
    """
    Load thermal propagated state from disk. Return None if the file is not found.

    Args:
        mol_list (:class:`MolList`): system information
        path (str): the path to load thermal state from. Should be an `npz` file.
    Returns: Loaded MpDm
    """
    try:
        logger.info(f"Try load from {path}")
        mpdm = MpDm.load(mol_list, path)
        logger.info(f"Init mpdm loaded: {mpdm}")
    except FileNotFoundError:
        logger.info(f"No file found in {path}")
        mpdm = None

    return mpdm