plotting stuff for project-nanotech

plot.py

import glob
import csv
import re
import collections
import ast
import numpy as np
from collections import defaultdict
import plotly.plotly as py
import plotly.graph_objs as go


def main():
    all_elements = import_csvs()
    averaged_elements = average_elements(all_elements)
    py.sign_in('tusing', 'redacted')
    plot_element_3d(averaged_elements)


def plot_element_3d(elements, keep=['freq', 'imp', 'phase']):

    data = []
    for element, tests in elements.items():
        trace = go.Scatter3d(
            x=tests[keep[0]],
            y=tests[keep[1]],
            z=tests[keep[2]],
            marker=dict(
                size=2,
            ),
            line=dict(
                width=1
            )
        )
        data.append(trace)

    layout = dict(
        autosize=True,
        title='Elements',
        scene=dict(
            xaxis=dict(
                gridcolor='rgb(255, 255, 255)',
                zerolinecolor='rgb(255, 255, 255)',
                showbackground=True,
                backgroundcolor='rgb(230, 230,230)',
                title = keep[0]
            ),
            yaxis=dict(
                gridcolor='rgb(255, 255, 255)',
                zerolinecolor='rgb(255, 255, 255)',
                showbackground=True,
                backgroundcolor='rgb(230, 230,230)',
                title = keep[1]
            ),
            zaxis=dict(
                gridcolor='rgb(255, 255, 255)',
                zerolinecolor='rgb(255, 255, 255)',
                showbackground=True,
                backgroundcolor='rgb(230, 230,230)',
                title = keep[2]
            ),
        ),
    )

    fig = dict(data=data, layout=layout)
    py.iplot(fig, filename='elements 3d' + ' '.join(keep))


def import_csvs(csv_query="current_testing_set/*",
                item_query='(\d+)_',
                keep={1: 'freq',
                      2: 'imp1der',
                      3: 'imp2der',
                      4: 'imp',
                      5: 'phase'}):

    csv_files = glob.glob(csv_query)
    all_elements = collections.defaultdict(lambda: defaultdict(list))

    for csv_file in csv_files:
        try:
            current_element_name = re.search(item_query, csv_file).group(1)
        except:
            continue
        with open(csv_file, newline='') as f:
            element_test = collections.defaultdict(list)
            reader = csv.reader(f, delimiter=';')
            next(reader)
            for row in reader:
                for k, v in keep.items():
                    element_test[v].append(ast.literal_eval(row[k]))
            for k, v in element_test.items():
                all_elements[current_element_name][k].append(v)
    return all_elements


def average_elements(all_elements):
    averaged_elements = collections.defaultdict(lambda: defaultdict(list))
    for element_name, element_tests in all_elements.items():
        for test, runs in element_tests.items():
            average_run = np.mean(np.array(runs), axis=0)
            averaged_elements[element_name][test] = average_run
    return averaged_elements

if __name__ == '__main__':
    main()