<commit_before>from ..learner import Learner
from ..utils import max_qvalue

class QLearning(Learner):
    '''Tabular Q-learning'''

    def __init__(self, env, policy, qfunction, learning_rate=0.1,
                 discount_factor=0.99, **kwargs):
        super(QLearning, self).__init__(env, **kwargs)
        self.policy = policy
        self.qfunction = qfunction
        self.learning_rate = learning_rate
        self.discount_factor = discount_factor

    def best_qvalue(self, state):
        return max_qvalue(state, self.env.actions(state), self.qfunction)

    ###########
    # Learner #
    ###########

    def episode(self):
        while not self.env.is_terminal():
            state = self.env.cur_state()
            action = self.policy.action(state)
            reward, next_state = self.env.do_action(action)
            best_qvalue = self.best_qvalue(next_state)
            target = reward + (self.discount_factor * best_qvalue)
            td_error = target - self.qfunction[state, action]
            self.qfunction[state, action] += self.learning_rate * td_error

class ApproximateQLearning(Learner):
    '''Q-learning with a function approximator'''

    def __init__(self, env, policy, qfunction, discount_factor=0.99, **kwargs):
        super(ApproximateQLearning, self).__init__(env, **kwargs)
        self.policy = policy
        self.qfunction = qfunction
        self.discount_factor = discount_factor

    def best_qvalue(self, state):
        return max_qvalue(state, self.env.actions(state), self.qfunction)

    ###########
    # Learner #
    ###########

    def episode(self):
        while not self.env.is_terminal():
            state = self.env.cur_state()
            action = self.policy.action(state)
            reward, next_state = self.env.do_action(action)
            best_qvalue = self.best_qvalue(next_state)
            update = reward + (self.discount_factor * best_qvalue)
            self.qfunction.update(state, update)
<commit_msg>Add experience replay<commit_after>import random
from ..learner import Learner
from ..utils import max_qvalue

class QLearning(Learner):
    '''Tabular Q-learning'''

    def __init__(self, env, policy, qfunction, learning_rate=0.1,
                 discount_factor=0.99, **kwargs):
        super(QLearning, self).__init__(env, **kwargs)
        self.policy = policy
        self.qfunction = qfunction
        self.learning_rate = learning_rate
        self.discount_factor = discount_factor

    def best_qvalue(self, state):
        return max_qvalue(state, self.env.actions(state), self.qfunction)

    ###########
    # Learner #
    ###########

    def episode(self):
        while not self.env.is_terminal():
            state = self.env.cur_state()
            action = self.policy.action(state)
            reward, next_state = self.env.do_action(action)
            best_qvalue = self.best_qvalue(next_state)
            target = reward + (self.discount_factor * best_qvalue)
            td_error = target - self.qfunction[state, action]
            self.qfunction[state, action] += self.learning_rate * td_error

class ApproximateQLearning(Learner):
    '''Q-learning with a function approximator'''

    def __init__(self, env, policy, qfunction, discount_factor=0.99,
                 experience_replay=True, **kwargs):
        super(ApproximateQLearning, self).__init__(env, **kwargs)
        self.policy = policy
        self.qfunction = qfunction
        self.discount_factor = discount_factor
        self.experience_replay = experience_replay
        if self.experience_replay:
            self.memory = set()

    def best_qvalue(self, state):
        return max_qvalue(state, self.env.actions(state), self.qfunction)

    ###########
    # Learner #
    ###########

    def episode(self):
        while not self.env.is_terminal():
            state = self.env.cur_state()
            action = self.policy.action(state)
            reward, next_state = self.env.do_action(action)
            if self.experience_replay:
                self.memory.add((state, action, reward, next_state))
                state, _, reward, next_state = random.choice(tuple(self.memory))
            best_qvalue = self.best_qvalue(next_state)
            update = reward + (self.discount_factor * best_qvalue)
            self.qfunction.update(state, update)
<|endoftext|>


<commit_before>from cauldron import cli
import sys
import re
from unittest.mock import patch

import cauldron
from cauldron import environ
from cauldron.cli import commander
from cauldron.cli.shell import CauldronShell
from cauldron.test.support import scaffolds
from cauldron.test.support.messages import Message
from cauldron.cli import parse
from cauldron.test.support import server

try:
  import readline
except ImportError:
  import pyreadline as readline

def run_remote_command(
        command: str,
        app=None,
        remote_connection: 'environ.RemoteConnection' = None,
        mock_send_request=None
) -> 'environ.Response':
    """ """

    name, args = parse.split_line(command)

    if not remote_connection:
        remote_connection = environ.RemoteConnection(
            url='fake-run-remote.command',
            active=True
        )

    app = app if app else server.create_test_app()

    def default_mock_send_request(
            endpoint: str,
            data: dict = None,
            method: str = None,
            **kwargs
    ):
        http_method = method.lower() if method else None
        func = server.post if data or http_method == 'post' else server.get
        result = func(
            app=app,
            endpoint=endpoint,
            data=data
        )
        return result.response

    side_effect = (
        mock_send_request
        if mock_send_request else
        default_mock_send_request
    )

    with patch(
            'cauldron.cli.sync.comm.send_request',
            side_effect=side_effect
    ):
        response = commander.execute(
            name=name,
            raw_args=args,
            remote_connection=remote_connection
        )
        response.thread.join()
        return response

def run_command(command: str) -> 'environ.Response':
    """

    :param command:
    :return:
    """

    cs = CauldronShell()
    cs.default(command)
    return cs.last_response

def create_project(
        tester: scaffolds.ResultsTest,
        name: str,
        path: str = None,
        forget: bool = True,
        confirm: bool = True,
        remote_connection: cli.CommandContext = None,
        **kwargs
) -> 'environ.Response':
    """

    :param tester:
    :param name:
    :param path:
    :param forget:
    :param confirm:
    :param remote_connection:
    :param kwargs:
    :return:
    """

    version = ''.join(['{}'.format(s) for s in sys.version_info])

    if path is None:
        path = tester.get_temp_path('project-{}-{}-'.format(name, version))

    args = [name, path]
    for key, value in kwargs.items():
        args.append('--{}="{}"'.format(key, value))

    if forget:
        args.append('--forget')

    args = ' '.join([a for a in args if a and len(a) > 0])

    r = commander.execute('create', args, remote_connection=remote_connection)
    if r.thread:
        r.thread.join()

    if confirm:
        tester.assertFalse(r.failed, Message(
            'support.create_project',
            'project should have been created',
            response=r
        ))
        tester.assertIsNotNone(cauldron.project.internal_project)

    return r

def open_project(
        tester: scaffolds.ResultsTest,
        path: str
) -> 'environ.Response':
    """

    :param tester:
    :param path:
    :return:
    """

    r = environ.Response()
    commander.execute('open', '{} --forget'.format(path), r)
    r.thread.join()
    return r

def autocomplete(command: str):
    """ """

    # On Linux/OSX the completer delims are retrieved from the readline module,
    # but the delims are different on Windows. So for testing consistency we
    # supply the Linux/OSX delims explicitly here in place of:
    # readline.get_completer_delims()

    completer_delims = ' \t\n`~!@#$%^&*()-=+[{]}\\|;:\'",<>/?'

    with patch('readline.get_line_buffer', return_value=command):
        cs = CauldronShell()
        line = command.lstrip()

        splits = re.split('[{}]{{1}}'.format(
            re.escape(completer_delims)),
            line
        )

        return cs.completedefault(
            text=splits[-1],
            line=line,
            begin_index=len(line) - len(splits[-1]),
            end_index=len(line) - 1
        )

def add_step(
        tester: 'scaffolds.ResultsTest',
        name: str = '',
        contents: str = '',
        position: str = None
):
    """

    :param tester:
    :param name:
    :param contents:
    :param position:
    :return:
    """

    cmd = [
        'steps add "{}"'.format(name),
        '--position="{}"'.format(position) if position else None
    ]

    r = run_command(' '.join([c for c in cmd if c]))
    step_path = r.data['step_path']

    with open(step_path, 'w+') as f:
        f.write(contents)

    return step_path
<commit_msg>Test Support<commit_after>from cauldron import cli
import sys
import re
import time
from unittest.mock import patch

import cauldron
from cauldron import environ
from cauldron.cli import commander
from cauldron.cli.shell import CauldronShell
from cauldron.test.support import scaffolds
from cauldron.test.support.messages import Message
from cauldron.cli import parse
from cauldron.test.support import server

try:
  import readline
except ImportError:
  import pyreadline as readline

def run_remote_command(
        command: str,
        app=None,
        remote_connection: 'environ.RemoteConnection' = None,
        mock_send_request=None
) -> 'environ.Response':
    """ """

    name, args = parse.split_line(command)

    if not remote_connection:
        remote_connection = environ.RemoteConnection(
            url='fake-run-remote.command',
            active=True
        )

    app = app if app else server.create_test_app()

    def default_mock_send_request(
            endpoint: str,
            data: dict = None,
            method: str = None,
            **kwargs
    ):
        http_method = method.lower() if method else None
        func = server.post if data or http_method == 'post' else server.get
        result = func(
            app=app,
            endpoint=endpoint,
            data=data
        )
        return result.response

    side_effect = (
        mock_send_request
        if mock_send_request else
        default_mock_send_request
    )

    with patch(
            'cauldron.cli.sync.comm.send_request',
            side_effect=side_effect
    ):
        response = commander.execute(
            name=name,
            raw_args=args,
            remote_connection=remote_connection
        )
        response.thread.join()
        return response

def run_command(command: str) -> 'environ.Response':
    """

    :param command:
    :return:
    """

    cs = CauldronShell()
    cs.default(command)
    return cs.last_response

def create_project(
        tester: scaffolds.ResultsTest,
        name: str,
        path: str = None,
        forget: bool = True,
        confirm: bool = True,
        remote_connection: cli.CommandContext = None,
        **kwargs
) -> 'environ.Response':
    """

    :param tester:
    :param name:
    :param path:
    :param forget:
    :param confirm:
    :param remote_connection:
    :param kwargs:
    :return:
    """

    version = ''.join(['{}'.format(s) for s in sys.version_info])

    if path is None:
        path = tester.get_temp_path('project-{}-{}-'.format(name, version))

    args = [name, path]
    for key, value in kwargs.items():
        args.append('--{}="{}"'.format(key, value))

    if forget:
        args.append('--forget')

    args = ' '.join([a for a in args if a and len(a) > 0])

    r = commander.execute('create', args, remote_connection=remote_connection)
    if r.thread:
        r.thread.join()

    # Prevent threading race conditions
    time.sleep(0.25)

    if confirm:
        tester.assertFalse(r.failed, Message(
            'support.create_project',
            ['project should have been created'],
            response=r
        ))
        tester.assertIsNotNone(cauldron.project.internal_project)

    return r

def open_project(
        tester: scaffolds.ResultsTest,
        path: str
) -> 'environ.Response':
    """

    :param tester:
    :param path:
    :return:
    """

    r = environ.Response()
    commander.execute('open', '{} --forget'.format(path), r)
    r.thread.join()
    return r

def autocomplete(command: str):
    """ """

    # On Linux/OSX the completer delims are retrieved from the readline module,
    # but the delims are different on Windows. So for testing consistency we
    # supply the Linux/OSX delims explicitly here in place of:
    # readline.get_completer_delims()

    completer_delims = ' \t\n`~!@#$%^&*()-=+[{]}\\|;:\'",<>/?'

    with patch('readline.get_line_buffer', return_value=command):
        cs = CauldronShell()
        line = command.lstrip()

        splits = re.split('[{}]{{1}}'.format(
            re.escape(completer_delims)),
            line
        )

        return cs.completedefault(
            text=splits[-1],
            line=line,
            begin_index=len(line) - len(splits[-1]),
            end_index=len(line) - 1
        )

def add_step(
        tester: 'scaffolds.ResultsTest',
        name: str = '',
        contents: str = '',
        position: str = None
):
    """

    :param tester:
    :param name:
    :param contents:
    :param position:
    :return:
    """

    cmd = [
        'steps add "{}"'.format(name),
        '--position="{}"'.format(position) if position else None
    ]

    r = run_command(' '.join([c for c in cmd if c]))
    step_path = r.data['step_path']

    with open(step_path, 'w+') as f:
        f.write(contents)

    return step_path
<|endoftext|>


<commit_before>import logging

import numpy as np

from .script import Script
from ..utils.errorvalue import ErrorValue

logger = logging.getLogger(__name__)
logger.setLevel(logging.DEBUG)

class Transmission(Script):
    """Measure the transmission of a sample

    Invocation: transmission(<samplename> [, <nimages> [, <countingtime [, <emptyname>]]])

    Arguments:
        <samplename>: the name of the sample. Can also be a list of strings if
            you want to measure multiple samples at once.
        <nimages>: number of images to expose. Integer, >2
        <countingtime>: counting time at each exposure
        <emptyname>: the sample to use for empty beam exposures

    """
    name = 'transmission'

    script = """
        #initialization of variables
        set('samplenames',_scriptargs[0])
        set('nimages',_scriptargs[1])
        set('exptime',_scriptargs[2])
        set('emptyname',_scriptargs[3])
        # initialization of instrument
        print('Initializing the instrument for transmission measurement')
        shutter('close')
        xray_power('low')
        beamstop('out')

        @startloop
            set('currentsample', samplenames.pop())
            print('Measuring transmission for sample',currentsample)
            print('Exposing dark current')
            exposemulti(exptime, nimages, _config['path']['prefixes']['tra'], 0.02, (currentsample, 'dark', nimages))
            print('Exposing empty beam')
            sample(emptyname)
            shutter('open')
            exposemulti(exptime, nimages, _config['path']['prefixes']['tra'], 0.02, (currentsample, 'empty', nimages))
            shutter('close')
            print('Exposing sample')
            sample(currentsample)
            shutter('open')
            exposemulti(exptime, nimages, _config['path']['prefixes']['tra'], 0.02, (currentsample, 'sample', nimages))
            shutter('close')
            goif('startloop', len(samplenames)>0)
        print('End of transmission measurements, moving beamstop into the beam')
        beamstop('in')
        end()
        """

    def execute(self, interpreter, arglist, instrument, namespace):
        self._instrument = instrument
        if isinstance(arglist[0], str):
            samplenames = [arglist[0]]
        else:
            samplenames = arglist[0]
        if len(arglist) < 2:
            nimages = self._instrument.config['transmission']['nimages']
        else:
            nimages = int(arglist[1])
        assert (nimages > 2)

        if len(arglist) < 3:
            exptime = self._instrument.config['transmission']['exptime']
        else:
            exptime = float(arglist[2])
        assert (exptime > 0)

        if len(arglist) < 4:
            emptyname = self._instrument.config['transmission']['empty_sample']
        else:
            emptyname = arglist[3]

        self._instrument_connections = [
            instrument.exposureanalyzer.connect('transmdata', self.on_transmdata),
        ]
        self._intensities = {}
        self._cannot_return_yet = True
        self._nsamples = len(samplenames)
        Script.execute(self, interpreter, (samplenames, nimages, exptime, emptyname), instrument, namespace)

    def on_transmdata(self, exposureanalyzer, prefix, fsn, data):
        logger.debug('Transmission data received: %s, %d, %s' % (prefix, fsn, data))
        samplename, what, nimages, I = data
        if samplename not in self._intensities:
            self._intensities[samplename] = {'dark': [], 'empty': [], 'sample': []}

        self._intensities[samplename][what].append(I)
        if len(self._intensities[samplename][what]) == nimages:
            self._intensities[samplename][what] = ErrorValue(
                np.mean(self._intensities[samplename][what]),
                np.std(self._intensities[samplename][what])
            )
            self.emit('message', 'I_%s for sample %s is: %s' % (
                what, samplename, self._intensities[samplename][what].tostring()))
            if what == 'sample':
                transm = ((self._intensities[samplename]['sample'] -
                           self._intensities[samplename]['dark']) /
                          (self._intensities[samplename]['empty'] -
                           self._intensities[samplename]['dark']))
                sam = self._instrument.samplestore.get_sample(samplename)
                sam.transmission = transm
                self._instrument.samplestore.set_sample(sam.title, sam)
                self.emit('message',
                          'Transmission value %s has been saved for sample %s.' % (transm.tostring(), sam.title))
                if self._nsamples == len(self._intensities):
                    del self._cannot_return_yet

    def cleanup(self):
        try:
            for c in self._instrument_connections:
                self._instrument.disconnect(c)
            del self._instrument_connections
        except AttributeError:
            pass
        return Script.cleanup(self)
<commit_msg>Added detail signal in transmission command<commit_after>import logging

import numpy as np

from .script import Script
from ..utils.errorvalue import ErrorValue

logger = logging.getLogger(__name__)
logger.setLevel(logging.DEBUG)

class Transmission(Script):
    """Measure the transmission of a sample

    Invocation: transmission(<samplename> [, <nimages> [, <countingtime [, <emptyname>]]])

    Arguments:
        <samplename>: the name of the sample. Can also be a list of strings if
            you want to measure multiple samples at once.
        <nimages>: number of images to expose. Integer, >2
        <countingtime>: counting time at each exposure
        <emptyname>: the sample to use for empty beam exposures

    """
    name = 'transmission'

    script = """
        #initialization of variables
        set('samplenames',_scriptargs[0])
        set('nimages',_scriptargs[1])
        set('exptime',_scriptargs[2])
        set('emptyname',_scriptargs[3])
        # initialization of instrument
        print('Initializing the instrument for transmission measurement')
        shutter('close')
        xray_power('low')
        beamstop('out')

        @startloop
            set('currentsample', samplenames.pop())
            print('Measuring transmission for sample',currentsample)
            print('Exposing dark current')
            exposemulti(exptime, nimages, _config['path']['prefixes']['tra'], 0.02, (currentsample, 'dark', nimages))
            print('Exposing empty beam')
            sample(emptyname)
            shutter('open')
            exposemulti(exptime, nimages, _config['path']['prefixes']['tra'], 0.02, (currentsample, 'empty', nimages))
            shutter('close')
            print('Exposing sample')
            sample(currentsample)
            shutter('open')
            exposemulti(exptime, nimages, _config['path']['prefixes']['tra'], 0.02, (currentsample, 'sample', nimages))
            shutter('close')
            goif('startloop', len(samplenames)>0)
        print('End of transmission measurements, moving beamstop into the beam')
        beamstop('in')
        end()
        """

    def execute(self, interpreter, arglist, instrument, namespace):
        self._instrument = instrument
        if isinstance(arglist[0], str):
            samplenames = [arglist[0]]
        else:
            samplenames = arglist[0]
        if len(arglist) < 2:
            nimages = self._instrument.config['transmission']['nimages']
        else:
            nimages = int(arglist[1])
        assert (nimages > 2)

        if len(arglist) < 3:
            exptime = self._instrument.config['transmission']['exptime']
        else:
            exptime = float(arglist[2])
        assert (exptime > 0)

        if len(arglist) < 4:
            emptyname = self._instrument.config['transmission']['empty_sample']
        else:
            emptyname = arglist[3]

        self._instrument_connections = [
            instrument.exposureanalyzer.connect('transmdata', self.on_transmdata),
        ]
        self._intensities = {}
        self._cannot_return_yet = True
        self._nsamples = len(samplenames)
        Script.execute(self, interpreter, (samplenames, nimages, exptime, emptyname), instrument, namespace)

    def on_transmdata(self, exposureanalyzer, prefix, fsn, data):
        logger.debug('Transmission data received: %s, %d, %s' % (prefix, fsn, data))
        samplename, what, nimages, I = data
        if samplename not in self._intensities:
            self._intensities[samplename] = {'dark': [], 'empty': [], 'sample': []}

        self._intensities[samplename][what].append(I)
        if len(self._intensities[samplename][what]) == nimages:
            self._intensities[samplename][what] = ErrorValue(
                np.mean(self._intensities[samplename][what]),
                np.std(self._intensities[samplename][what])
            )
            self.emit('message', 'I_%s for sample %s is: %s' % (
                what, samplename, self._intensities[samplename][what].tostring()))
            self.emit('detail', (what, samplename, self._intensities[samplename][what]))
            if what == 'sample':
                transm = ((self._intensities[samplename]['sample'] -
                           self._intensities[samplename]['dark']) /
                          (self._intensities[samplename]['empty'] -
                           self._intensities[samplename]['dark']))
                sam = self._instrument.samplestore.get_sample(samplename)
                sam.transmission = transm
                self._instrument.samplestore.set_sample(sam.title, sam)
                self.emit('message',
                          'Transmission value %s has been saved for sample %s.' % (transm.tostring(), sam.title))
                self.emit('detail', ('transmission', samplename, transm))
                if self._nsamples == len(self._intensities):
                    del self._cannot_return_yet

    def cleanup(self):
        try:
            for c in self._instrument_connections:
                self._instrument.disconnect(c)
            del self._instrument_connections
        except AttributeError:
            pass
        return Script.cleanup(self)
<|endoftext|>


<commit_before>"""
This is a template script - an example of how a CEA script should be set up.
NOTE: ADD YOUR SCRIPT'S DOCUMENTATION HERE (what, why, include literature references)
"""
from __future__ import division
from __future__ import print_function

import math
import os

import numpy as np
import osmnx as ox
from geopandas import GeoDataFrame as Gdf
import pandas as pd

import cea.config
import cea.inputlocator
from cea.datamanagement.databases_verification import COLUMNS_ZONE_TYPOLOGY
from cea.demand import constants
from cea.utilities.dbf import dataframe_to_dbf
from cea.utilities.standardize_coordinates import get_projected_coordinate_system, get_geographic_coordinate_system

__author__ = "Jimeno Fonseca"
__copyright__ = "Copyright 2019, Architecture and Building Systems - ETH Zurich"
__credits__ = ["Jimeno Fonseca", "Reynold Mok"]
__license__ = "MIT"
__version__ = "0.1"
__maintainer__ = "Daren Thomas"
__email__ = "cea@arch.ethz.ch"
__status__ = "Production"

def check_if_float(x):
    y = 0.0
    try:
        y = float(x)
    except:
        if ',' in x:
            y = float(x.split(',')[-1])
    return y

def clean_attributes(shapefile, buildings_height, buildings_floors, buildings_height_below_ground,
                     buildings_floors_below_ground, key):
    # local variables
    no_buildings = shapefile.shape[0]
    list_of_columns = shapefile.columns
    if buildings_height is None and buildings_floors is None:
        print('Warning! you have not indicated a height or number of floors above ground for the buildings, '
              'we are reverting to data stored in Open Street Maps (It might not be accurate at all),'
              'if we do not find data in OSM for a particular building, we get the median in the surroundings, '
              'if we do not get any data we assume 4 floors per building')

        # Check which attributes the OSM has, Sometimes it does not have any and indicate the data source
        if 'building:levels' not in list_of_columns:
            shapefile['building:levels'] = [3] * no_buildings
            shapefile['REFERENCE'] = "CEA - assumption"
        elif pd.isnull(shapefile['building:levels']).all():
            shapefile['building:levels'] = [3] * no_buildings
            shapefile['REFERENCE'] = "CEA - assumption"
        else:
            shapefile['REFERENCE'] = ["OSM - median values of all buildings" if x is np.nan else "OSM - as it is" for x in
                                        shapefile['building:levels']]
        if 'roof:levels' not in list_of_columns:
            shapefile['roof:levels'] = [1] * no_buildings

        # get the median from the area:
        data_osm_floors1 = shapefile['building:levels'].fillna(0)
        data_osm_floors2 = shapefile['roof:levels'].fillna(0)
        data_floors_sum = [x + y for x, y in zip([check_if_float(x) for x in data_osm_floors1], [check_if_float(x) for x in data_osm_floors2])]
        data_floors_sum_with_nan = [np.nan if x < 1.0 else x for x in data_floors_sum]
        data_osm_floors_joined = int(
            math.ceil(np.nanmedian(data_floors_sum_with_nan)))  # median so we get close to the worse case
        shapefile["floors_ag"] = [int(x) if x is not np.nan else data_osm_floors_joined for x in
                                  data_floors_sum_with_nan]
        shapefile["height_ag"] = shapefile["floors_ag"] * constants.H_F
    else:
        shapefile['REFERENCE'] = "User - assumption"
        if buildings_height is None and buildings_floors is not None:
            shapefile["floors_ag"] = [buildings_floors] * no_buildings
            shapefile["height_ag"] = shapefile["floors_ag"] * constants.H_F
        elif buildings_height is not None and buildings_floors is None:
            shapefile["height_ag"] = [buildings_height] * no_buildings
            shapefile["floors_ag"] = [int(math.floor(x)) for x in shapefile["height_ag"] / constants.H_F]
        else:  # both are not none
            shapefile["height_ag"] = [buildings_height] * no_buildings
            shapefile["floors_ag"] = [buildings_floors] * no_buildings

    # add fields for floorsa and height below ground
    shapefile["height_bg"] = [buildings_height_below_ground] * no_buildings
    shapefile["floors_bg"] = [buildings_floors_below_ground] * no_buildings

    # add description
    if "description" in list_of_columns:
        shapefile["description"] = shapefile['description']
    elif 'addr:housename' in list_of_columns:
        shapefile["description"] = shapefile['addr:housename']
    elif 'amenity' in list_of_columns:
        shapefile["description"] = shapefile['amenity']
    else:
        shapefile["description"] = [np.nan] * no_buildings

    shapefile["category"] = shapefile['building']
    shapefile["Name"] = [key + str(x + 1000) for x in
                         range(no_buildings)]  # start in a big number to avoid potential confusion

    result = shapefile[
        ["Name", "height_ag", "floors_ag", "height_bg", "floors_bg", "description", "category", "geometry",
         "REFERENCE"]]

    result.reset_index(inplace=True, drop=True)
    shapefile.reset_index(inplace=True, drop=True)

    return result, shapefile

def zone_helper(locator, config):
    """
    This script gets a polygon and calculates the zone.shp and the occupancy.dbf and age.dbf inputs files for CEA
    :param locator:
    :param config:
    :return:
    """
    # local variables:
    poly = Gdf.from_file(locator.get_site_polygon())
    buildings_height = config.zone_helper.height_ag
    buildings_floors = config.zone_helper.floors_ag
    buildings_height_below_ground = config.zone_helper.height_bg
    buildings_floors_below_ground = config.zone_helper.floors_bg
    occupancy_type = config.zone_helper.occupancy_type
    year_construction = config.zone_helper.year_construction
    zone_output_path = locator.get_zone_geometry()
    typology_output_path = locator.get_building_typology()

    # ensure folders exist
    locator.ensure_parent_folder_exists(zone_output_path)
    locator.ensure_parent_folder_exists(typology_output_path)

    # get zone.shp file and save in folder location
    zone_df = polygon_to_zone(buildings_floors, buildings_floors_below_ground, buildings_height,
                              buildings_height_below_ground,
                              poly, zone_output_path)

    # USE_A zone.shp file contents to get the contents of occupancy.dbf and age.dbf
    calculate_typology_file(locator, zone_df, year_construction, occupancy_type, typology_output_path)

def calc_category(standard_DB, year_array):

    def category_assignment(year):
        return (standard_DB[(standard_DB['YEAR_START'] <= year) & (standard_DB['YEAR_END'] >= year)].STANDARD.values[0])

    category = np.vectorize(category_assignment)(year_array)
    return category

def calculate_typology_file(locator, zone_df, year_construction, occupancy_type, occupancy_output_path):
    """
    This script fills in the occupancy.dbf file with one occupancy type
    :param zone_df:
    :param occupancy_type:
    :param occupancy_output_path:
    :return:
    """
    #calculate construction year
    typology_df = calculate_age(zone_df, year_construction)

    #calculate the most likely construction standard
    standard_database = pd.read_excel(locator.get_database_construction_standards(), sheet_name='STANDARD_DEFINITION')
    typology_df['STANDARD'] = calc_category(standard_database, typology_df['YEAR'].values)

    #Calculate the most likely use type
    typology_df['1ST_USE'] = 'MULTI_RES'
    typology_df['1ST_USE_R'] = 1.0
    typology_df['2ND_USE'] = "NONE"
    typology_df['2ND_USE_R'] = 0.0
    typology_df['3RD_USE'] = "NONE"
    typology_df['3RD_USE_R'] = 0.0
    if occupancy_type == "Get it from open street maps":
        no_buildings = typology_df.shape[0]
        for index in range(no_buildings):
            typology_df.loc[index, "USE_A_R"] = 1.0
            if zone_df.loc[index, "category"] == "yes":
                typology_df.loc[index, "USE_A"] = "MULTI_RES"
                typology_df.loc[index, "REFERENCE"] = "CEA - assumption"
            elif zone_df.loc[index, "category"] == "residential" or zone_df.loc[index, "category"] == "apartments":
                typology_df.loc[index, "USE_A"] = "MULTI_RES"
                typology_df.loc[index, "REFERENCE"] = "OSM - as it is"
            elif zone_df.loc[index, "category"] == "commercial" or zone_df.loc[index, "category"] == "civic":
                typology_df.loc[index, "USE_A"] = "OFFICE"
                typology_df.loc[index, "REFERENCE"] = "OSM - as it is"
            elif zone_df.loc[index, "category"] == "school":
                typology_df.loc[index, "USE_A"] = "SCHOOL"
                typology_df.loc[index, "REFERENCE"] = "OSM - as it is"
            elif zone_df.loc[index, "category"] == "garage" or zone_df.loc[index, "category"] == "garages" or zone_df.loc[index, "category"] == "warehouse":
                typology_df.loc[index, "USE_A"] = "PARKING"
                typology_df.loc[index, "REFERENCE"] = "OSM - as it is"
            elif zone_df.loc[index, "category"] == "house" or zone_df.loc[index, "category"] == "terrace" or zone_df.loc[index, "category"] == "detached":
                typology_df.loc[index, "USE_A"] = "SINGLE_RES"
                typology_df.loc[index, "REFERENCE"] = "OSM - as it is"
            elif zone_df.loc[index, "category"] == "retail":
                typology_df.loc[index, "USE_A"] = "RETAIL"
                typology_df.loc[index, "REFERENCE"] = "OSM - as it is"
            elif zone_df.loc[index, "category"] == "industrial":
                typology_df.loc[index, "USE_A"] = "INDUSTRIAL"
                typology_df.loc[index, "REFERENCE"] = "OSM - as it is"
            elif zone_df.loc[index, "category"] == "warehouse":
                typology_df.loc[index, "USE_A"] = "INDUSTRIAL"
                typology_df.loc[index, "REFERENCE"] = "OSM - as it is"
            else:
                typology_df.loc[index, "USE_A"] = "MULTI_RES"
                typology_df.loc[index, "REFERENCE"] = "CEA - assumption"

    fields = COLUMNS_ZONE_TYPOLOGY
    dataframe_to_dbf(typology_df[fields+['REFERENCE']], occupancy_output_path)

def calculate_age(zone_df, year_construction):
    """
    This script fills in the age.dbf file with one year of construction
    :param zone_df:
    :param year_construction:
    :param age_output_path:
    :return:
    """
    if year_construction is None:
        print('Warning! you have not indicated a year of construction for the buildings, '
              'we are reverting to data stored in Open Street Maps (It might not be accurate at all),'
              'if we do not find data in OSM for a particular building, we get the median in the surroundings, '
              'if we do not get any data we assume all buildings being constructed in the year 2000')
        list_of_columns = zone_df.columns
        if "start_date" not in list_of_columns:  # this field describes the construction year of buildings
            zone_df["start_date"] = 2000
            zone_df['REFERENCE'] = "CEA - assumption"
        else:
            zone_df['REFERENCE'] = ["OSM - median" if x is np.nan else "OSM - as it is" for x in zone_df['start_date']]

        data_floors_sum_with_nan = [np.nan if x is np.nan else int(x) for x in zone_df['start_date']]
        data_osm_floors_joined = int(math.ceil(np.nanmedian(data_floors_sum_with_nan)))  # median so we get close to the worse case
        zone_df["YEAR"] = [int(x) if x is not np.nan else data_osm_floors_joined for x in data_floors_sum_with_nan]
    else:
        zone_df['REFERENCE'] = "CEA - assumption"

    return zone_df

def polygon_to_zone(buildings_floors, buildings_floors_below_ground, buildings_height, buildings_height_below_ground,
                    poly, shapefile_out_path):
    poly = poly.to_crs(get_geographic_coordinate_system())
    lon = poly.geometry[0].centroid.coords.xy[0][0]
    lat = poly.geometry[0].centroid.coords.xy[1][0]
    # get footprints of all the district
    poly = ox.footprints.create_footprints_gdf(polygon=poly['geometry'].values[0])
    # clean attributes of height, name and number of floors
    result, result_allfields = clean_attributes(poly, buildings_height, buildings_floors, buildings_height_below_ground,
                              buildings_floors_below_ground, key="B")
    result = result.to_crs(get_projected_coordinate_system(float(lat), float(lon)))
    # save to shapefile
    result.to_file(shapefile_out_path)

    return result_allfields

def main(config):
    """
    This script gets a polygon and calculates the zone.shp and the occupancy.dbf and age.dbf inputs files for CEA
    """
    assert os.path.exists(config.scenario), 'Scenario not found: %s' % config.scenario
    locator = cea.inputlocator.InputLocator(config.scenario)

    zone_helper(locator, config)

if __name__ == '__main__':
    main(cea.config.Configuration())
<commit_msg>hotfix issue 2547<commit_after>"""
This is a template script - an example of how a CEA script should be set up.
NOTE: ADD YOUR SCRIPT'S DOCUMENTATION HERE (what, why, include literature references)
"""
from __future__ import division
from __future__ import print_function

import math
import os

import numpy as np
import osmnx as ox
from geopandas import GeoDataFrame as Gdf
import pandas as pd

import cea.config
import cea.inputlocator
from cea.datamanagement.databases_verification import COLUMNS_ZONE_TYPOLOGY
from cea.demand import constants
from cea.utilities.dbf import dataframe_to_dbf
from cea.utilities.standardize_coordinates import get_projected_coordinate_system, get_geographic_coordinate_system

__author__ = "Jimeno Fonseca"
__copyright__ = "Copyright 2019, Architecture and Building Systems - ETH Zurich"
__credits__ = ["Jimeno Fonseca", "Reynold Mok"]
__license__ = "MIT"
__version__ = "0.1"
__maintainer__ = "Daren Thomas"
__email__ = "cea@arch.ethz.ch"
__status__ = "Production"

def check_if_float(x):
    y = 0.0
    try:
        y = float(x)
    except:
        if ',' in x:
            y = float(x.split(',')[-1])
    return y

def clean_attributes(shapefile, buildings_height, buildings_floors, buildings_height_below_ground,
                     buildings_floors_below_ground, key):
    # local variables
    no_buildings = shapefile.shape[0]
    list_of_columns = shapefile.columns
    if buildings_height is None and buildings_floors is None:
        print('Warning! you have not indicated a height or number of floors above ground for the buildings, '
              'we are reverting to data stored in Open Street Maps (It might not be accurate at all),'
              'if we do not find data in OSM for a particular building, we get the median in the surroundings, '
              'if we do not get any data we assume 4 floors per building')

        # Check which attributes the OSM has, Sometimes it does not have any and indicate the data source
        if 'building:levels' not in list_of_columns:
            shapefile['building:levels'] = [3] * no_buildings
            shapefile['REFERENCE'] = "CEA - assumption"
        elif pd.isnull(shapefile['building:levels']).all():
            shapefile['building:levels'] = [3] * no_buildings
            shapefile['REFERENCE'] = "CEA - assumption"
        else:
            shapefile['REFERENCE'] = ["OSM - median values of all buildings" if x is np.nan else "OSM - as it is" for x in
                                        shapefile['building:levels']]
        if 'roof:levels' not in list_of_columns:
            shapefile['roof:levels'] = [1] * no_buildings

        # get the median from the area:
        data_osm_floors1 = shapefile['building:levels'].fillna(0)
        data_osm_floors2 = shapefile['roof:levels'].fillna(0)
        data_floors_sum = [x + y for x, y in zip([check_if_float(x) for x in data_osm_floors1], [check_if_float(x) for x in data_osm_floors2])]
        data_floors_sum_with_nan = [np.nan if x < 1.0 else x for x in data_floors_sum]
        data_osm_floors_joined = int(
            math.ceil(np.nanmedian(data_floors_sum_with_nan)))  # median so we get close to the worse case
        shapefile["floors_ag"] = [int(x) if x is not np.nan else data_osm_floors_joined for x in
                                  data_floors_sum_with_nan]
        shapefile["height_ag"] = shapefile["floors_ag"] * constants.H_F
    else:
        shapefile['REFERENCE'] = "User - assumption"
        if buildings_height is None and buildings_floors is not None:
            shapefile["floors_ag"] = [buildings_floors] * no_buildings
            shapefile["height_ag"] = shapefile["floors_ag"] * constants.H_F
        elif buildings_height is not None and buildings_floors is None:
            shapefile["height_ag"] = [buildings_height] * no_buildings
            shapefile["floors_ag"] = [int(math.floor(x)) for x in shapefile["height_ag"] / constants.H_F]
        else:  # both are not none
            shapefile["height_ag"] = [buildings_height] * no_buildings
            shapefile["floors_ag"] = [buildings_floors] * no_buildings

    # add fields for floorsa and height below ground
    shapefile["height_bg"] = [buildings_height_below_ground] * no_buildings
    shapefile["floors_bg"] = [buildings_floors_below_ground] * no_buildings

    # add description
    if "description" in list_of_columns:
        shapefile["description"] = shapefile['description']
    elif 'addr:housename' in list_of_columns:
        shapefile["description"] = shapefile['addr:housename']
    elif 'amenity' in list_of_columns:
        shapefile["description"] = shapefile['amenity']
    else:
        shapefile["description"] = [np.nan] * no_buildings

    shapefile["category"] = shapefile['building']
    shapefile["Name"] = [key + str(x + 1000) for x in
                         range(no_buildings)]  # start in a big number to avoid potential confusion

    result = shapefile[
        ["Name", "height_ag", "floors_ag", "height_bg", "floors_bg", "description", "category", "geometry",
         "REFERENCE"]]

    result.reset_index(inplace=True, drop=True)
    shapefile.reset_index(inplace=True, drop=True)

    return result, shapefile

def zone_helper(locator, config):
    """
    This script gets a polygon and calculates the zone.shp and the occupancy.dbf and age.dbf inputs files for CEA
    :param locator:
    :param config:
    :return:
    """
    # local variables:
    poly = Gdf.from_file(locator.get_site_polygon())
    buildings_height = config.zone_helper.height_ag
    buildings_floors = config.zone_helper.floors_ag
    buildings_height_below_ground = config.zone_helper.height_bg
    buildings_floors_below_ground = config.zone_helper.floors_bg
    occupancy_type = config.zone_helper.occupancy_type
    year_construction = config.zone_helper.year_construction
    zone_output_path = locator.get_zone_geometry()
    typology_output_path = locator.get_building_typology()

    # ensure folders exist
    locator.ensure_parent_folder_exists(zone_output_path)
    locator.ensure_parent_folder_exists(typology_output_path)

    # get zone.shp file and save in folder location
    zone_df = polygon_to_zone(buildings_floors, buildings_floors_below_ground, buildings_height,
                              buildings_height_below_ground,
                              poly, zone_output_path)

    # USE_A zone.shp file contents to get the contents of occupancy.dbf and age.dbf
    calculate_typology_file(locator, zone_df, year_construction, occupancy_type, typology_output_path)

def calc_category(standard_DB, year_array):

    def category_assignment(year):
        return (standard_DB[(standard_DB['YEAR_START'] <= year) & (standard_DB['YEAR_END'] >= year)].STANDARD.values[0])

    category = np.vectorize(category_assignment)(year_array)
    return category

def calculate_typology_file(locator, zone_df, year_construction, occupancy_type, occupancy_output_path):
    """
    This script fills in the occupancy.dbf file with one occupancy type
    :param zone_df:
    :param occupancy_type:
    :param occupancy_output_path:
    :return:
    """
    #calculate construction year
    typology_df = calculate_age(zone_df, year_construction)

    #calculate the most likely construction standard
    standard_database = pd.read_excel(locator.get_database_construction_standards(), sheet_name='STANDARD_DEFINITION')
    typology_df['STANDARD'] = calc_category(standard_database, typology_df['YEAR'].values)

    #Calculate the most likely use type
    typology_df['1ST_USE'] = 'MULTI_RES'
    typology_df['1ST_USE_R'] = 1.0
    typology_df['2ND_USE'] = "NONE"
    typology_df['2ND_USE_R'] = 0.0
    typology_df['3RD_USE'] = "NONE"
    typology_df['3RD_USE_R'] = 0.0
    if occupancy_type == "Get it from open street maps":
        no_buildings = typology_df.shape[0]
        for index in range(no_buildings):
            typology_df.loc[index, "USE_A_R"] = 1.0
            if zone_df.loc[index, "category"] == "yes":
                typology_df.loc[index, "USE_A"] = "MULTI_RES"
                typology_df.loc[index, "REFERENCE"] = "CEA - assumption"
            elif zone_df.loc[index, "category"] == "residential" or zone_df.loc[index, "category"] == "apartments":
                typology_df.loc[index, "USE_A"] = "MULTI_RES"
                typology_df.loc[index, "REFERENCE"] = "OSM - as it is"
            elif zone_df.loc[index, "category"] == "commercial" or zone_df.loc[index, "category"] == "civic":
                typology_df.loc[index, "USE_A"] = "OFFICE"
                typology_df.loc[index, "REFERENCE"] = "OSM - as it is"
            elif zone_df.loc[index, "category"] == "school":
                typology_df.loc[index, "USE_A"] = "SCHOOL"
                typology_df.loc[index, "REFERENCE"] = "OSM - as it is"
            elif zone_df.loc[index, "category"] == "garage" or zone_df.loc[index, "category"] == "garages" or zone_df.loc[index, "category"] == "warehouse":
                typology_df.loc[index, "USE_A"] = "PARKING"
                typology_df.loc[index, "REFERENCE"] = "OSM - as it is"
            elif zone_df.loc[index, "category"] == "house" or zone_df.loc[index, "category"] == "terrace" or zone_df.loc[index, "category"] == "detached":
                typology_df.loc[index, "USE_A"] = "SINGLE_RES"
                typology_df.loc[index, "REFERENCE"] = "OSM - as it is"
            elif zone_df.loc[index, "category"] == "retail":
                typology_df.loc[index, "USE_A"] = "RETAIL"
                typology_df.loc[index, "REFERENCE"] = "OSM - as it is"
            elif zone_df.loc[index, "category"] == "industrial":
                typology_df.loc[index, "USE_A"] = "INDUSTRIAL"
                typology_df.loc[index, "REFERENCE"] = "OSM - as it is"
            elif zone_df.loc[index, "category"] == "warehouse":
                typology_df.loc[index, "USE_A"] = "INDUSTRIAL"
                typology_df.loc[index, "REFERENCE"] = "OSM - as it is"
            else:
                typology_df.loc[index, "USE_A"] = "MULTI_RES"
                typology_df.loc[index, "REFERENCE"] = "CEA - assumption"

    fields = COLUMNS_ZONE_TYPOLOGY
    dataframe_to_dbf(typology_df[fields+['REFERENCE']], occupancy_output_path)

def cast_year(year):
    try:
        if year.split('-'):
            return year.split('-')[0]
        else:
            return year
    except:
        raise Exception("We could not cast the value",year)

def calculate_age(zone_df, year_construction):
    """
    This script fills in the age.dbf file with one year of construction
    :param zone_df:
    :param year_construction:
    :param age_output_path:
    :return:
    """
    if year_construction is None:
        print('Warning! you have not indicated a year of construction for the buildings, '
              'we are reverting to data stored in Open Street Maps (It might not be accurate at all),'
              'if we do not find data in OSM for a particular building, we get the median in the surroundings, '
              'if we do not get any data we assume all buildings being constructed in the year 2000')
        list_of_columns = zone_df.columns
        if "start_date" not in list_of_columns:  # this field describes the construction year of buildings
            zone_df["start_date"] = 2000
            zone_df['REFERENCE'] = "CEA - assumption"
        else:
            zone_df['REFERENCE'] = ["OSM - median" if x is np.nan else "OSM - as it is" for x in zone_df['start_date']]

        data_floors_sum_with_nan = [np.nan if x is np.nan else int(cast_year(x)) for x in zone_df['start_date']]
        data_osm_floors_joined = int(math.ceil(np.nanmedian(data_floors_sum_with_nan)))  # median so we get close to the worse case
        zone_df["YEAR"] = [int(x) if x is not np.nan else data_osm_floors_joined for x in data_floors_sum_with_nan]
    else:
        zone_df['REFERENCE'] = "CEA - assumption"

    return zone_df

def polygon_to_zone(buildings_floors, buildings_floors_below_ground, buildings_height, buildings_height_below_ground,
                    poly, shapefile_out_path):
    poly = poly.to_crs(get_geographic_coordinate_system())
    lon = poly.geometry[0].centroid.coords.xy[0][0]
    lat = poly.geometry[0].centroid.coords.xy[1][0]
    # get footprints of all the district
    poly = ox.footprints.create_footprints_gdf(polygon=poly['geometry'].values[0])
    # clean attributes of height, name and number of floors
    result, result_allfields = clean_attributes(poly, buildings_height, buildings_floors, buildings_height_below_ground,
                              buildings_floors_below_ground, key="B")
    result = result.to_crs(get_projected_coordinate_system(float(lat), float(lon)))
    # save to shapefile
    result.to_file(shapefile_out_path)

    return result_allfields

def main(config):
    """
    This script gets a polygon and calculates the zone.shp and the occupancy.dbf and age.dbf inputs files for CEA
    """
    assert os.path.exists(config.scenario), 'Scenario not found: %s' % config.scenario
    locator = cea.inputlocator.InputLocator(config.scenario)

    zone_helper(locator, config)

if __name__ == '__main__':
    main(cea.config.Configuration())
<|endoftext|>


<commit_before># ~/ceam/ceam/modules/disease_models.py

from datetime import timedelta

from ceam import config
from ceam.framework.state_machine import Transition, State, TransitionSet
from ceam.framework.disease import DiseaseModel, DiseaseState, ExcessMortalityState, IncidenceRateTransition, ProportionTransition
# from ceam.gbd_data.gbd_ms_functions import get_angina_proportions, get_heart_failure_incidence_draws

def heart_disease_factory():
    module = DiseaseModel('ihd')

    healthy = State('healthy', key='ihd')

    cause_of_heart_failure_me_id = 2414

    # Calculate an adjusted disability weight for the acute heart attack phase that
    # accounts for the fact that our timestep is longer than the phase length
    # TODO: This doesn't account for the fact that our timestep is longer than 28 days
    timestep = config.getfloat('simulation_parameters', 'time_step')
    weight = 0.43*(2/timestep) + 0.07*(28/timestep)
    heart_attack = ExcessMortalityState('heart_attack', disability_weight=weight, dwell_time=timedelta(days=28), modelable_entity_id=1814, prevalence_meid=1814)

    #
    mild_heart_failure = ExcessMortalityState('mild_heart_failure', disability_weight=0.04, modelable_entity_id=2412, prevalence_meid=1821)
    moderate_heart_failure = ExcessMortalityState('moderate_heart_failure', disability_weight=0.07, modelable_entity_id=2412, prevalence_meid=1822)
    severe_heart_failure = ExcessMortalityState('severe_heart_failure', disability_weight=0.18, modelable_entity_id=2412, prevalence_meid=1823)

    asymptomatic_angina = ExcessMortalityState('asymptomatic_angina', disability_weight=0.0, modelable_entity_id=1817, prevalence_meid=3102)
    mild_angina = ExcessMortalityState('mild_angina', disability_weight=0.03, modelable_entity_id=1817, prevalence_meid=1818)
    moderate_angina = ExcessMortalityState('moderate_angina', disability_weight=0.08, modelable_entity_id=1817, prevalence_meid=1819)
    severe_angina = ExcessMortalityState('severe_angina', disability_weight=0.17, modelable_entity_id=1817, prevalence_meid=1820)
    asymptomatic_ihd = ExcessMortalityState('asymptomatic_ihd', disability_weight=0.00, modelable_entity_id=3233, prevalence_meid=3233)

    heart_attack_transition = IncidenceRateTransition(heart_attack, 'heart_attack', modelable_entity_id=1814)
    healthy.transition_set.append(heart_attack_transition)

    heart_failure_buckets = TransitionSet(allow_null_transition=False, key="heart_failure_split")
    heart_failure_buckets.extend([
        ProportionTransition(mild_heart_failure, proportion=0.182074),
        ProportionTransition(moderate_heart_failure, proportion=0.149771),
        ProportionTransition(severe_heart_failure, proportion=0.402838),
        ])

    angina_buckets = TransitionSet(allow_null_transition=False, key="angina_split")
    angina_buckets.extend([
        ProportionTransition(asymptomatic_angina, proportion=0.304553),
        ProportionTransition(mild_angina, proportion=0.239594),
        ProportionTransition(moderate_angina, proportion=0.126273),
        ProportionTransition(severe_angina, proportion=0.32958),
        ])
    healthy.transition_set.append(IncidenceRateTransition(angina_buckets, 'non_mi_angina', modelable_entity_id=1817))

    heart_attack.transition_set.allow_null_transition=False

    hf_prop_df = normalize_for_simulation(load_data_from_cache(get_heart_failure_incidence_draws(location_id, year_start, year_end, cause_of_heart_failure_me_id)))
    angina_prop_df = normalize_for_simulation(load_data_from_cache(get_angina_proportions(start_year, end_year)))
    asympt_prop_df = get_post_mi_asympt_ihd_proportion(hf_prop_df, angina_prop_df)

    # post-mi transitions
    # TODO: Figure out if we can pass in me_id here to get incidence for the correct cause of heart failure
    # TODO: Figure out how to make asymptomatic ihd be equal to whatever is left after people get heart failure and angina
    heart_attack.transition_set.append(ProportionTransition(heart_failure_buckets, proportion=hf_prop_df))
    heart_attack.transition_set.append(ProportionTransition(angina_buckets, proportion=angina_prop_df))
    heart_attack.transition_set.append(ProportionTransition(asymptomatic_ihd, proportion=asympt_prop_df)

    mild_heart_failure.transition_set.append(heart_attack_transition)
    moderate_heart_failure.transition_set.append(heart_attack_transition)
    severe_heart_failure.transition_set.append(heart_attack_transition)
    asymptomatic_angina.transition_set.append(heart_attack_transition)
    mild_angina.transition_set.append(heart_attack_transition)
    moderate_angina.transition_set.append(heart_attack_transition)
    severe_angina.transition_set.append(heart_attack_transition)
    asymptomatic_ihd.transition_set.append(heart_attack_transition)

    module.states.extend([healthy, heart_attack, mild_heart_failure, moderate_heart_failure, severe_heart_failure, asymptomatic_angina, mild_angina, moderate_angina, severe_angina, asymptomatic_ihd])
    return module

def stroke_factory():
    module = DiseaseModel('hemorrhagic_stroke')

    healthy = State('healthy', key='hemorrhagic_stroke')
    # TODO: disability weight for stroke
    hemorrhagic_stroke = ExcessMortalityState('hemorrhagic_stroke', disability_weight=0.32, dwell_time=timedelta(days=28), modelable_entity_id=9311)
    ischemic_stroke = ExcessMortalityState('ischemic_stroke', disability_weight=0.32, dwell_time=timedelta(days=28), modelable_entity_id=9310)
    chronic_stroke = ExcessMortalityState('chronic_stroke', disability_weight=0.32, modelable_entity_id=9312)

    hemorrhagic_transition = IncidenceRateTransition(hemorrhagic_stroke, 'hemorrhagic_stroke', modelable_entity_id=9311)
    ischemic_transition = IncidenceRateTransition(ischemic_stroke, 'ischemic_stroke', modelable_entity_id=9310)
    healthy.transition_set.extend([hemorrhagic_transition, ischemic_transition])

    hemorrhagic_stroke.transition_set.append(Transition(chronic_stroke))
    ischemic_stroke.transition_set.append(Transition(chronic_stroke))

    chronic_stroke.transition_set.extend([hemorrhagic_transition, ischemic_transition])

    module.states.extend([healthy, hemorrhagic_stroke, ischemic_stroke, chronic_stroke])

    return module

# End.
<commit_msg>only test left is to do a full run<commit_after># ~/ceam/ceam/modules/disease_models.py

from datetime import timedelta

from ceam import config
from ceam.framework.state_machine import Transition, State, TransitionSet
from ceam.framework.disease import DiseaseModel, DiseaseState, ExcessMortalityState, IncidenceRateTransition, ProportionTransition
# from ceam.gbd_data.gbd_ms_functions import get_angina_proportions, get_heart_failure_incidence_draws

def heart_disease_factory():
    module = DiseaseModel('ihd')

    healthy = State('healthy', key='ihd')

    cause_of_heart_failure_me_id = 2414

    # Calculate an adjusted disability weight for the acute heart attack phase that
    # accounts for the fact that our timestep is longer than the phase length
    # TODO: This doesn't account for the fact that our timestep is longer than 28 days
    timestep = config.getfloat('simulation_parameters', 'time_step')
    weight = 0.43*(2/timestep) + 0.07*(28/timestep)
    heart_attack = ExcessMortalityState('heart_attack', disability_weight=weight, dwell_time=timedelta(days=28), modelable_entity_id=1814, prevalence_meid=1814)

    #
    mild_heart_failure = ExcessMortalityState('mild_heart_failure', disability_weight=0.04, modelable_entity_id=2412, prevalence_meid=1821)
    moderate_heart_failure = ExcessMortalityState('moderate_heart_failure', disability_weight=0.07, modelable_entity_id=2412, prevalence_meid=1822)
    severe_heart_failure = ExcessMortalityState('severe_heart_failure', disability_weight=0.18, modelable_entity_id=2412, prevalence_meid=1823)

    asymptomatic_angina = ExcessMortalityState('asymptomatic_angina', disability_weight=0.0, modelable_entity_id=1817, prevalence_meid=3102)
    mild_angina = ExcessMortalityState('mild_angina', disability_weight=0.03, modelable_entity_id=1817, prevalence_meid=1818)
    moderate_angina = ExcessMortalityState('moderate_angina', disability_weight=0.08, modelable_entity_id=1817, prevalence_meid=1819)
    severe_angina = ExcessMortalityState('severe_angina', disability_weight=0.17, modelable_entity_id=1817, prevalence_meid=1820)
    asymptomatic_ihd = ExcessMortalityState('asymptomatic_ihd', disability_weight=0.00, modelable_entity_id=3233, prevalence_meid=3233)

    heart_attack_transition = IncidenceRateTransition(heart_attack, 'heart_attack', modelable_entity_id=1814)
    healthy.transition_set.append(heart_attack_transition)

    heart_failure_buckets = TransitionSet(allow_null_transition=False, key="heart_failure_split")
    heart_failure_buckets.extend([
        ProportionTransition(mild_heart_failure, proportion=0.182074),
        ProportionTransition(moderate_heart_failure, proportion=0.149771),
        ProportionTransition(severe_heart_failure, proportion=0.402838),
        ])

    angina_buckets = TransitionSet(allow_null_transition=False, key="angina_split")
    angina_buckets.extend([
        ProportionTransition(asymptomatic_angina, proportion=0.304553),
        ProportionTransition(mild_angina, proportion=0.239594),
        ProportionTransition(moderate_angina, proportion=0.126273),
        ProportionTransition(severe_angina, proportion=0.32958),
        ])
    healthy.transition_set.append(IncidenceRateTransition(angina_buckets, 'non_mi_angina', modelable_entity_id=1817))

    heart_attack.transition_set.allow_null_transition=False

    hf_prop_df = normalize_for_simulation(load_data_from_cache(get_heart_failure_incidence_draws(location_id, year_start, year_end, cause_of_heart_failure_me_id)))
    angina_prop_df = normalize_for_simulation(load_data_from_cache(get_angina_proportions(start_year, end_year)))
    asympt_prop_df = get_post_mi_asympt_ihd_proportion(hf_prop_df, angina_prop_df)

    # post-mi transitions
    # TODO: Figure out if we can pass in me_id here to get incidence for the correct cause of heart failure
    # TODO: Figure out how to make asymptomatic ihd be equal to whatever is left after people get heart failure and angina
    heart_attack.transition_set.append(ProportionTransition(heart_failure_buckets, proportion=hf_prop_df))
    heart_attack.transition_set.append(ProportionTransition(angina_buckets, proportion=angina_prop_df))
    heart_attack.transition_set.append(ProportionTransition(asymptomatic_ihd, proportion=asympt_prop_df))

    mild_heart_failure.transition_set.append(heart_attack_transition)
    moderate_heart_failure.transition_set.append(heart_attack_transition)
    severe_heart_failure.transition_set.append(heart_attack_transition)
    asymptomatic_angina.transition_set.append(heart_attack_transition)
    mild_angina.transition_set.append(heart_attack_transition)
    moderate_angina.transition_set.append(heart_attack_transition)
    severe_angina.transition_set.append(heart_attack_transition)
    asymptomatic_ihd.transition_set.append(heart_attack_transition)

    module.states.extend([healthy, heart_attack, mild_heart_failure, moderate_heart_failure, severe_heart_failure, asymptomatic_angina, mild_angina, moderate_angina, severe_angina, asymptomatic_ihd])
    return module

def stroke_factory():
    module = DiseaseModel('hemorrhagic_stroke')

    healthy = State('healthy', key='hemorrhagic_stroke')
    # TODO: disability weight for stroke
    hemorrhagic_stroke = ExcessMortalityState('hemorrhagic_stroke', disability_weight=0.32, dwell_time=timedelta(days=28), modelable_entity_id=9311)
    ischemic_stroke = ExcessMortalityState('ischemic_stroke', disability_weight=0.32, dwell_time=timedelta(days=28), modelable_entity_id=9310)
    chronic_stroke = ExcessMortalityState('chronic_stroke', disability_weight=0.32, modelable_entity_id=9312)

    hemorrhagic_transition = IncidenceRateTransition(hemorrhagic_stroke, 'hemorrhagic_stroke', modelable_entity_id=9311)
    ischemic_transition = IncidenceRateTransition(ischemic_stroke, 'ischemic_stroke', modelable_entity_id=9310)
    healthy.transition_set.extend([hemorrhagic_transition, ischemic_transition])

    hemorrhagic_stroke.transition_set.append(Transition(chronic_stroke))
    ischemic_stroke.transition_set.append(Transition(chronic_stroke))

    chronic_stroke.transition_set.extend([hemorrhagic_transition, ischemic_transition])

    module.states.extend([healthy, hemorrhagic_stroke, ischemic_stroke, chronic_stroke])

    return module

# End.
<|endoftext|>


<commit_before>"""pec csv-type data files"""
import os
from dateutil.parser import parse
import logging
import warnings
import numpy as np

import pandas as pd

from cellpy.readers.core import FileID, Cell, humanize_bytes
from cellpy.parameters.internal_settings import get_headers_normal
from cellpy.readers.instruments.mixin import Loader

pec_headers_normal = dict()

pec_headers_normal["step_index_txt"] = "Step"
pec_headers_normal["cycle_index_txt"] = "Cycle"
pec_headers_normal["test_time_txt"] = "Total_Time_Seconds"
pec_headers_normal["step_time_txt"] = "Step_Time_Seconds"
pec_headers_normal["datetime_txt"] = "Real_Time"
pec_headers_normal["voltage_txt"] = "Voltage_mV"
pec_headers_normal["current_txt"] = "Current_mA"
pec_headers_normal["charge_capacity_txt"] = "Charge_Capacity_mAh"
pec_headers_normal["discharge_capacity_txt"] = "Discharge_Capacity_mAh"
pec_headers_normal["charge_energy_txt"] = "Charge_Capacity_mWh"
pec_headers_normal["discharge_energy_txt"] = "Discharge_Capacity_mWh"
pec_headers_normal["internal_resistance_txt"] = "Internal_Resistance_1_mOhm"
pec_headers_normal["test_id_txt"] = "Test"

# TODO: better reading of first part of the file (comments and headers)
#  1. find the units

#  2. find user-defined variables
#  3. find units
# Hei
# Og hei til deg!
# og hei

class PECLoader(Loader):
    """Main loading class"""

    def __init__(self):
        self.headers_normal = (
            get_headers_normal()
        )  # should consider to move this to the Loader class
        self.current_chunk = 0  # use this to set chunks to load
        self.pec_data = None
        self.pec_log = None
        self.pec_settings = None
        self.number_of_header_lines = 32  # Number of header lines is not constant
        self.cellpy_headers = (
            get_headers_normal()
        )  # should consider to move this to the Loader class

    #@staticmethod
    #def _get_pec_units():
    #    pec_units = dict()
    #    pec_units["voltage"] = 0.001  # V
    #    pec_units["current"] = 0.001  # A
    #    pec_units["charge"] = 0.001  # Ah
    #    pec_units["mass"] = 0.001  # g
    #    pec_units["energy"] = 0.001  # Wh

    #    return pec_units

    def _get_pec_units(self, filename):  # Fetches units from a csv file
        if not os.path.isfile(file_name):
            return None

        # Mapping prefixes to values
        prefix = {
            'µ': 10 ** -6,
            'm': 10 ** -3,
            '': 1
        }

        # Adding the non-variable units to the return value
        pec_units = {
            'charge': 0.001,  # Ah
            'mass': 0.001,  # g
            'energy': 0.001  # Wh
        }

        # Searching values for variable units
        header = ['Voltage (V)', 'Current (A)']

        data = pd.read_csv(filename, skiprows=self.number_of_header_lines)

        for item in data.keys():
            for unit in header:
                x = unit.find('(') - len(unit)
                if unit[:x + 1] in item:
                    y = item[x].replace('(', '')
                    # Adding units to return value
                    if header.index(unit) == 0:
                        pec_units['voltage'] = prefix.get(y)
                    elif header.index(unit) == 1:
                        pec_units['current'] = prefix.get(y)

        return pec_units

    @staticmethod
    def get_raw_units():
        """Include the settings for the units used by the instrument.

        The units are defined w.r.t. the SI units ('unit-fractions'; currently only units that are multiples of
        Si units can be used). For example, for current defined in mA, the value for the
        current unit-fraction will be 0.001.

        Returns: dictionary containing the unit-fractions for current, charge, and mass

        """

        raw_units = dict()
        raw_units["voltage"] = 1.0  # V
        raw_units["current"] = 1.0  # A
        raw_units["charge"] = 1.0  # Ah
        raw_units["mass"] = 0.001  # g
        raw_units["energy"] = 1.0  # Wh

        return raw_units

    def get_raw_limits(self):
        """Include the settings for how to decide what kind of step you are examining here.

        The raw limits are 'epsilons' used to check if the current and/or voltage is stable (for example
        for galvanostatic steps, one would expect that the current is stable (constant) and non-zero).
        It is expected that different instruments (with different resolution etc.) have different
        'epsilons'.

        Returns: the raw limits (dict)

        """
        warnings.warn("raw limits have not been subject for testing yet")
        raw_limits = dict()
        raw_limits["current_hard"] = 0.1  # There is a bug in PEC
        raw_limits["current_soft"] = 1.0
        raw_limits["stable_current_hard"] = 2.0
        raw_limits["stable_current_soft"] = 4.0
        raw_limits["stable_voltage_hard"] = 2.0
        raw_limits["stable_voltage_soft"] = 4.0
        raw_limits["stable_charge_hard"] = 2.0
        raw_limits["stable_charge_soft"] = 5.0
        raw_limits["ir_change"] = 0.00001
        return raw_limits

    def loader(self, file_name, bad_steps=None, **kwargs):
        new_tests = []
        if not os.path.isfile(file_name):
            self.logger.info("Missing file_\n   %s" % file_name)
            return None

        filesize = os.path.getsize(file_name)
        hfilesize = humanize_bytes(filesize)
        txt = "Filesize: %i (%s)" % (filesize, hfilesize)
        logging.debug(txt)

        data = Cell()
        fid = FileID(file_name)

        # div parameters and information (probably load this last)
        test_no = 1
        data.cell_no = test_no
        data.loaded_from = file_name

        # some overall prms
        data.channel_index = None
        data.channel_number = None
        data.creator = None
        data.item_ID = None
        data.schedule_file_name = None
        data.test_ID = None
        data.test_name = None
        data.raw_data_files.append(fid)

        # --------- read raw-data (normal-data) -------------------------

        self._load_pec_data(file_name, bad_steps)
        data.start_datetime = self.pec_settings["start_time"]
        length_of_test = self.pec_data.shape[0]
        logging.debug(f"length of test: {length_of_test}")

        logging.debug("renaming columns")
        self._rename_headers()
        self._convert_units()

        data.raw = self.pec_data

        data.raw_data_files_length.append(length_of_test)
        new_tests.append(data)

        return new_tests

    def _load_pec_data(self, file_name, bad_steps):
        number_of_header_lines = self.number_of_header_lines

        # ----------------- reading the data ---------------------
        df = pd.read_csv(file_name, skiprows=number_of_header_lines)

        # get rid of unnamed columns
        df = df.loc[:, ~df.columns.str.contains("^Unnamed")]

        # get rid of spaces, parenthesis, and the deg-sign
        new_column_headers = {
            c: c.replace(" ", "_")
            .replace("(", "")
            .replace(")", "")
            .replace("°", "")
            .replace(r"%", "pct")
            for c in df.columns
        }
        df.rename(columns=new_column_headers, inplace=True)

        # add missing columns
        df.insert(0, self.headers_normal.data_point_txt, range(len(df)))
        df[self.headers_normal.sub_step_index_txt] = 0
        df[self.headers_normal.sub_step_time_txt] = 0

        self.pec_data = df

        # ----------------  reading the parameters ---------------
        with open(file_name, "r") as ofile:
            counter = 0
            lines = []
            for line in ofile:
                counter += 1
                if counter > number_of_header_lines:
                    break
                lines.append(line)
        self._extract_variables(lines)

    def _extract_variables(self, lines):
        header_comments = dict()
        comment_loop = False
        for line_number, line in enumerate(lines):

            if line.startswith("#"):
                if not comment_loop:
                    comment_loop = True
                else:
                    comment_loop = False

            else:
                if not comment_loop:
                    parts = line.split(",")
                    variable = parts[0].strip()
                    variable = variable.strip(":")
                    variable = variable.replace(" ", "_")
                    try:
                        value = parts[1].strip()
                    except IndexError:
                        value = None

                    if not value:
                        value = np.nan
                    header_comments[variable] = value
        logging.debug(" Headers Dict ")
        logging.debug(header_comments)

        headers = dict()

        start_time = parse(header_comments["Start_Time"])
        end_time = parse(header_comments["End_Time"])

        headers["start_time"] = start_time
        headers["end_time"] = end_time
        headers["test_regime_name"] = header_comments["TestRegime_Name"]

        self.pec_settings = headers

    def _rename_headers(self):
        logging.debug("Trying to rename the columns")
        # logging.debug("Current columns:")
        # logging.debug(self.pec_data.columns)
        # logging.debug("Rename to:")
        # logging.debug(self.headers_normal)

        for key in pec_headers_normal:
            self._rename_header(key, pec_headers_normal[key])

        # logging.debug("New cols:")
        # logging.debug(self.pec_data.columns)

    def _convert_units(self):
        logging.debug("Trying to convert all data into correct units")
        logging.debug("- dtypes")
        self.pec_data[self.headers_normal.datetime_txt] = pd.to_datetime(
            self.pec_data[self.headers_normal.datetime_txt]
        )

        self.pec_data["Position_Start_Time"] = pd.to_datetime(
            self.pec_data["Position_Start_Time"]
        )

        self.pec_data["Rack"] = self.pec_data["Rack"].astype("category")

        logging.debug("- cellpy units")
        pec_units = self._get_pec_units()
        raw_units = self.get_raw_units()

        _v = pec_units["voltage"] / raw_units["voltage"]
        _i = pec_units["current"] / raw_units["current"]
        _c = pec_units["charge"] / raw_units["charge"]
        _w = pec_units["energy"] / raw_units["energy"]

        v_txt = self.headers_normal.voltage_txt
        i_txt = self.headers_normal.current_txt

        self.pec_data[v_txt] *= _v
        self.pec_data[i_txt] *= _i

        self.pec_data[self.headers_normal.charge_capacity_txt] *= _c
        self.pec_data[self.headers_normal.discharge_capacity_txt] *= _c
        self.pec_data[self.headers_normal.charge_energy_txt] *= _w
        self.pec_data[self.headers_normal.discharge_energy_txt] *= _w

    def _rename_header(self, h_old, h_new):
        try:
            self.pec_data.rename(
                columns={h_new: self.cellpy_headers[h_old]}, inplace=True
            )
        except KeyError as e:
            logging.info(f"Problem during conversion to cellpy-format ({e})")

    def find_header_length(self, filename):
        if not os.path.isfile(filename):
            return None

        header = open(filename, 'r')
        skiprows = 0

        # Counting number of lines until the last line appears
        for line in header:
            skiprows += 1
            if line == "#END RESULTS CHECK\n":  # this is the last line of header
                break
        header.close()

        return skiprows

if __name__ == "__main__":
    pass
<commit_msg>Some optimizations to get_pec_units<commit_after>"""pec csv-type data files"""
import os
from dateutil.parser import parse
import logging
import warnings
import numpy as np

import pandas as pd

from cellpy.readers.core import FileID, Cell, humanize_bytes
from cellpy.parameters.internal_settings import get_headers_normal
from cellpy.readers.instruments.mixin import Loader

pec_headers_normal = dict()

pec_headers_normal["step_index_txt"] = "Step"
pec_headers_normal["cycle_index_txt"] = "Cycle"
pec_headers_normal["test_time_txt"] = "Total_Time_Seconds"
pec_headers_normal["step_time_txt"] = "Step_Time_Seconds"
pec_headers_normal["datetime_txt"] = "Real_Time"
pec_headers_normal["voltage_txt"] = "Voltage_mV"
pec_headers_normal["current_txt"] = "Current_mA"
pec_headers_normal["charge_capacity_txt"] = "Charge_Capacity_mAh"
pec_headers_normal["discharge_capacity_txt"] = "Discharge_Capacity_mAh"
pec_headers_normal["charge_energy_txt"] = "Charge_Capacity_mWh"
pec_headers_normal["discharge_energy_txt"] = "Discharge_Capacity_mWh"
pec_headers_normal["internal_resistance_txt"] = "Internal_Resistance_1_mOhm"
pec_headers_normal["test_id_txt"] = "Test"

# TODO: better reading of first part of the file (comments and headers)
#  1. find the units

#  2. find user-defined variables
#  3. find units
# Hei
# Og hei til deg!
# og hei

class PECLoader(Loader):
    """Main loading class"""

    def __init__(self):
        self.headers_normal = (
            get_headers_normal()
        )  # should consider to move this to the Loader class
        self.current_chunk = 0  # use this to set chunks to load
        self.pec_data = None
        self.pec_log = None
        self.pec_settings = None
        self.number_of_header_lines = 32  # Number of header lines is not constant
        self.cellpy_headers = (
            get_headers_normal()
        )  # should consider to move this to the Loader class

    #@staticmethod
    #def _get_pec_units():
    #    pec_units = dict()
    #    pec_units["voltage"] = 0.001  # V
    #    pec_units["current"] = 0.001  # A
    #    pec_units["charge"] = 0.001  # Ah
    #    pec_units["mass"] = 0.001  # g
    #    pec_units["energy"] = 0.001  # Wh

    #    return pec_units

    def _get_pec_units(self, filename):  # Fetches units from a csv file
        if not os.path.isfile(filename):
            return None

        # Mapping prefixes to values
        prefix = {
            'µ': 10 ** -6,
            'm': 10 ** -3,
            '': 1
        }

        # Adding the non-variable units to the return value
        pec_units = {
            'charge': 0.001,  # Ah
            'mass': 0.001,  # g
            'energy': 0.001  # Wh
        }

        # Searching values for variable units
        header = ['Voltage (V)', 'Current (A)']

        data = pd.read_csv(filename, skiprows=self.number_of_header_lines)

        for item in data.keys():
            for unit in header:
                x = unit.find('(') - len(unit)
                if unit[:x + 1] in item:
                    y = item[x].replace('(', '')
                    # Adding units to return value
                    if header.index(unit) == 0:
                        pec_units['voltage'] = prefix.get(y)
                    elif header.index(unit) == 1:
                        pec_units['current'] = prefix.get(y)

        return pec_units

    @staticmethod
    def get_raw_units():
        """Include the settings for the units used by the instrument.

        The units are defined w.r.t. the SI units ('unit-fractions'; currently only units that are multiples of
        Si units can be used). For example, for current defined in mA, the value for the
        current unit-fraction will be 0.001.

        Returns: dictionary containing the unit-fractions for current, charge, and mass

        """

        raw_units = dict()
        raw_units["voltage"] = 1.0  # V
        raw_units["current"] = 1.0  # A
        raw_units["charge"] = 1.0  # Ah
        raw_units["mass"] = 0.001  # g
        raw_units["energy"] = 1.0  # Wh

        return raw_units

    def get_raw_limits(self):
        """Include the settings for how to decide what kind of step you are examining here.

        The raw limits are 'epsilons' used to check if the current and/or voltage is stable (for example
        for galvanostatic steps, one would expect that the current is stable (constant) and non-zero).
        It is expected that different instruments (with different resolution etc.) have different
        'epsilons'.

        Returns: the raw limits (dict)

        """
        warnings.warn("raw limits have not been subject for testing yet")
        raw_limits = dict()
        raw_limits["current_hard"] = 0.1  # There is a bug in PEC
        raw_limits["current_soft"] = 1.0
        raw_limits["stable_current_hard"] = 2.0
        raw_limits["stable_current_soft"] = 4.0
        raw_limits["stable_voltage_hard"] = 2.0
        raw_limits["stable_voltage_soft"] = 4.0
        raw_limits["stable_charge_hard"] = 2.0
        raw_limits["stable_charge_soft"] = 5.0
        raw_limits["ir_change"] = 0.00001
        return raw_limits

    def loader(self, file_name, bad_steps=None, **kwargs):
        new_tests = []
        if not os.path.isfile(file_name):
            self.logger.info("Missing file_\n   %s" % file_name)
            return None

        filesize = os.path.getsize(file_name)
        hfilesize = humanize_bytes(filesize)
        txt = "Filesize: %i (%s)" % (filesize, hfilesize)
        logging.debug(txt)

        data = Cell()
        fid = FileID(file_name)

        # div parameters and information (probably load this last)
        test_no = 1
        data.cell_no = test_no
        data.loaded_from = file_name

        # some overall prms
        data.channel_index = None
        data.channel_number = None
        data.creator = None
        data.item_ID = None
        data.schedule_file_name = None
        data.test_ID = None
        data.test_name = None
        data.raw_data_files.append(fid)

        # --------- read raw-data (normal-data) -------------------------

        self._load_pec_data(file_name, bad_steps)
        data.start_datetime = self.pec_settings["start_time"]
        length_of_test = self.pec_data.shape[0]
        logging.debug(f"length of test: {length_of_test}")

        logging.debug("renaming columns")
        self._rename_headers()
        self._convert_units()

        data.raw = self.pec_data

        data.raw_data_files_length.append(length_of_test)
        new_tests.append(data)

        return new_tests

    def _load_pec_data(self, file_name, bad_steps):
        number_of_header_lines = self.number_of_header_lines

        # ----------------- reading the data ---------------------
        df = pd.read_csv(file_name, skiprows=number_of_header_lines)

        # get rid of unnamed columns
        df = df.loc[:, ~df.columns.str.contains("^Unnamed")]

        # get rid of spaces, parenthesis, and the deg-sign
        new_column_headers = {
            c: c.replace(" ", "_")
            .replace("(", "")
            .replace(")", "")
            .replace("°", "")
            .replace(r"%", "pct")
            for c in df.columns
        }
        df.rename(columns=new_column_headers, inplace=True)

        # add missing columns
        df.insert(0, self.headers_normal.data_point_txt, range(len(df)))
        df[self.headers_normal.sub_step_index_txt] = 0
        df[self.headers_normal.sub_step_time_txt] = 0

        self.pec_data = df

        # ----------------  reading the parameters ---------------
        with open(file_name, "r") as ofile:
            counter = 0
            lines = []
            for line in ofile:
                counter += 1
                if counter > number_of_header_lines:
                    break
                lines.append(line)
        self._extract_variables(lines)

    def _extract_variables(self, lines):
        header_comments = dict()
        comment_loop = False
        for line_number, line in enumerate(lines):

            if line.startswith("#"):
                if not comment_loop:
                    comment_loop = True
                else:
                    comment_loop = False

            else:
                if not comment_loop:
                    parts = line.split(",")
                    variable = parts[0].strip()
                    variable = variable.strip(":")
                    variable = variable.replace(" ", "_")
                    try:
                        value = parts[1].strip()
                    except IndexError:
                        value = None

                    if not value:
                        value = np.nan
                    header_comments[variable] = value
        logging.debug(" Headers Dict ")
        logging.debug(header_comments)

        headers = dict()

        start_time = parse(header_comments["Start_Time"])
        end_time = parse(header_comments["End_Time"])

        headers["start_time"] = start_time
        headers["end_time"] = end_time
        headers["test_regime_name"] = header_comments["TestRegime_Name"]

        self.pec_settings = headers

    def _rename_headers(self):
        logging.debug("Trying to rename the columns")
        # logging.debug("Current columns:")
        # logging.debug(self.pec_data.columns)
        # logging.debug("Rename to:")
        # logging.debug(self.headers_normal)

        for key in pec_headers_normal:
            self._rename_header(key, pec_headers_normal[key])

        # logging.debug("New cols:")
        # logging.debug(self.pec_data.columns)

    def _convert_units(self):
        logging.debug("Trying to convert all data into correct units")
        logging.debug("- dtypes")
        self.pec_data[self.headers_normal.datetime_txt] = pd.to_datetime(
            self.pec_data[self.headers_normal.datetime_txt]
        )

        self.pec_data["Position_Start_Time"] = pd.to_datetime(
            self.pec_data["Position_Start_Time"]
        )

        self.pec_data["Rack"] = self.pec_data["Rack"].astype("category")

        logging.debug("- cellpy units")
        pec_units = self._get_pec_units()
        raw_units = self.get_raw_units()

        _v = pec_units["voltage"] / raw_units["voltage"]
        _i = pec_units["current"] / raw_units["current"]
        _c = pec_units["charge"] / raw_units["charge"]
        _w = pec_units["energy"] / raw_units["energy"]

        v_txt = self.headers_normal.voltage_txt
        i_txt = self.headers_normal.current_txt

        self.pec_data[v_txt] *= _v
        self.pec_data[i_txt] *= _i

        self.pec_data[self.headers_normal.charge_capacity_txt] *= _c
        self.pec_data[self.headers_normal.discharge_capacity_txt] *= _c
        self.pec_data[self.headers_normal.charge_energy_txt] *= _w
        self.pec_data[self.headers_normal.discharge_energy_txt] *= _w

    def _rename_header(self, h_old, h_new):
        try:
            self.pec_data.rename(
                columns={h_new: self.cellpy_headers[h_old]}, inplace=True
            )
        except KeyError as e:
            logging.info(f"Problem during conversion to cellpy-format ({e})")

    def find_header_length(self, filename):
        if not os.path.isfile(filename):
            return None

        header = open(filename, 'r')
        skiprows = 0

        # Counting number of lines until the last line appears
        for line in header:
            skiprows += 1
            if line == "#END RESULTS CHECK\n":  # this is the last line of header
                break
        header.close()

        return skiprows

if __name__ == "__main__":
    pass
<|endoftext|>