<commit_before>import utility as ut
import glob
import re
import csv
import os
import numpy as np
import pandas as pd
import scipy.optimize as sop
import matplotlib.pyplot as plt
from functools import partial

# These regular expressions are used to extract information from the names
# of the data files
h_patt = '([+-]?\d+(?:\.\d+)?)Oe'
t_patt = '(\d+(?:\.\d+)?)K'
p_patt = '([+-]?\d+(?:\.\d+)?dBm)'
n_patt = '(\d+(?:\.\d+)?).txt'
f_patt = '(\d+(?:\.\d+)?)MHz'
pd_patt = 'P(\d+(?:\.\d+)?)D(\d+(?:\.\d+)?)'
defNSpatt = '.* ' + t_patt + ' ' + n_patt
defVNAoutpatt = '(.* 100).* ' + p_patt + ' *'
defpatt = '.* ' + h_patt + '.* ' + t_patt + ' .*.txt'
defFSFSpatt = '.* ' + h_patt + '.* ' + t_patt + ' .* ' + f_patt + '.*.txt'
defPulsepatt = '.* ' + t_patt + '.* ' + f_patt + '* ' + h_patt + '.*.txt'
defPulseFSpatt = '.* ' + f_patt + '*.txt'
defPulseoutpatt = '(.*Sweep).* (\d+Ave) ([+-]?\d+(?:\.\d+)?dBm)*'
defPSoutpatt = '(.*us) .* (\d+Ave) ([+-]?\d+(?:\.\d+)?dBm)*'
defechopatt = '.* ' + h_patt + '.* ' + t_patt + ' ' + pd_patt + '.* ' + n_patt
filepatt = '^K  ([+-]?\\d+(?:\\.\\d+)?)*'

# This section is for data gathered from loop-gap resonators using the VNA

# Names of the columns in an exported dataset
colnames = ["Magnetic Field", "Base", "B-error", "Peak Amplitude", "A-error",
            "FWHH", "W-error", "Central Frequency", "Frequency Shift",
            "F-error", "R-Squared", "Q-Value", "Q-error" "Minimum",
            "Temperature"]

# Functions for VNA fitting


def xp_to_quad(data, mod=1):
    """Converts VNA data from amplitude (x) and phase (p) to in-phase (i) and
    quadrature (q).

    'data' is a data array with elements [frequency (f), x, p];
    'mod' is 1 (default) to convert to [f, q, i], 3 to convert to [f, i, q],
    5 to convert to [f, -x, p], 0 to not convert and just pass [f, x, p].

    Returns one of the above arrays."""
    f = data[0]
    x = data[1]
    p = data[2]
    q = (-x * np.cos(p / 180 * np.pi) / (x.max() - x.min()) + x.mean()) * mod
    i = 180 / np.pi * x * np.sin(p / 180 * np.pi)

    def outsw(x):
        return {
            3: [f, np.sign(mod) * i, q],
            5: [f, -x, p],
            0: [f, x, p]
        }.get(x, [f, q, i])
    return np.array(outsw(np.abs(mod)))


def test_vna_xp(path, dir, skew, full, mod):
    """Loads a VNA dataset and fits it and its converted counterpart to
    Lorentzians, plotting the result to two separate subplots.

    The fit function is:
    y(x) = b + a * (w / 2)**2 / ((x - f)**2 + (w / 2)**2)

    b is the background offset, a is the peak amplitude, w is the peak width,
    and f is the peak frequency.

    'path' is a string corresponding to the desired datafile;
    'dir' is the data directory;
    'skew' is 1 to subtract a linear background, and 0 to use no slope;
    'full' is 1 to fit the whole dataset, and 0 to only fit between the
    maxima on either side of the peak;
    for details on 'mod' see 'xp_to_quad'."""
    name = glob.glob(os.path.join(dir, path))[0]
    data = np.loadtxt(name, skiprows=1).transpose()
    xdata = xp_to_quad(data, mod)
    f, axarr = plt.subplots(2, sharex=True, figsize=(8, 6))
    ut.plot_lor_fit(data, skew, full, axarr[0])
    ut.plot_lor_fit(xdata, skew, 0, axarr[1])


def load_vna_data(path, dir, mod=0, patt=defNSpatt):
    """Loads all datasets from field sweeps using the VNA corresponding to a
    given string.

    'path' is a string corresponding to the desired datafile;
    'dir' is the data directory;
    'mod' controls whether to convert the amplitude and phase data into
    in-phase and quadrature data, see documentation for
    'xp_to_quad' for full details (default:0, no conversion);
    'patt' is a regular expression pattern to extract temperature and sequence
    information from the filename.

    Returns a list with: a list of datasets; an array of experimental fields;
    an array of experimental temperatures; one of the matching filenames."""
    data, names, _ = ut.get_data_exclude_names(path + '*', dir, 'Min.txt', 1)

    def matchnum(name):
        [temp, num] = re.match(patt, name).groups()
        return [float(temp), int(num)]
    [temps, nums] = np.array(list(map(matchnum, names))).transpose()

    def fields1(path):
        with open(path) as f:
            first_line = f.readline()
        return float(re.match(filepatt, first_line).groups()[0])
    fields = np.array([fields1(nam) for nam in names])
    sortarr = np.argsort(nums)
    [data, names, fields, temps] = [np.array(data)[sortarr], np.array(names)[
        sortarr], fields[sortarr], temps[sortarr]]
    if mod != 0:
        data = np.array([xp_to_quad(dat, mod) for dat in data])
    return [data, fields, temps, os.path.split(names[0])[1]]
<commit_msg>Allow for loading VNA data from zipfiles, and exclude 'Fit' files<commit_after>import utility as ut
import glob
import re
import csv
import os
import numpy as np
import pandas as pd
import scipy.optimize as sop
import matplotlib.pyplot as plt
from functools import partial

# These regular expressions are used to extract information from the names
# of the data files
h_patt = '([+-]?\d+(?:\.\d+)?)Oe'
t_patt = '(\d+(?:\.\d+)?)K'
p_patt = '([+-]?\d+(?:\.\d+)?dBm)'
n_patt = '(\d+(?:\.\d+)?).txt'
f_patt = '(\d+(?:\.\d+)?)MHz'
pd_patt = 'P(\d+(?:\.\d+)?)D(\d+(?:\.\d+)?)'
defNSpatt = '.* ' + t_patt + ' ' + n_patt
defVNAoutpatt = '(.* 100).* ' + p_patt + ' *'
defpatt = '.* ' + h_patt + '.* ' + t_patt + ' .*.txt'
defFSFSpatt = '.* ' + h_patt + '.* ' + t_patt + ' .* ' + f_patt + '.*.txt'
defPulsepatt = '.* ' + t_patt + '.* ' + f_patt + '* ' + h_patt + '.*.txt'
defPulseFSpatt = '.* ' + f_patt + '*.txt'
defPulseoutpatt = '(.*Sweep).* (\d+Ave) ([+-]?\d+(?:\.\d+)?dBm)*'
defPSoutpatt = '(.*us) .* (\d+Ave) ([+-]?\d+(?:\.\d+)?dBm)*'
defechopatt = '.* ' + h_patt + '.* ' + t_patt + ' ' + pd_patt + '.* ' + n_patt
filepatt = '^K  ([+-]?\\d+(?:\\.\\d+)?)*'

# This section is for data gathered from loop-gap resonators using the VNA

# Names of the columns in an exported dataset
colnames = ["Magnetic Field", "Base", "B-error", "Peak Amplitude", "A-error",
            "FWHH", "W-error", "Central Frequency", "Frequency Shift",
            "F-error", "R-Squared", "Q-Value", "Q-error" "Minimum",
            "Temperature"]

# Functions for VNA fitting


def xp_to_quad(data, mod=1):
    """Converts VNA data from amplitude (x) and phase (p) to in-phase (i) and
    quadrature (q).

    'data' is a data array with elements [frequency (f), x, p];
    'mod' is 1 (default) to convert to [f, q, i], 3 to convert to [f, i, q],
    5 to convert to [f, -x, p], 0 to not convert and just pass [f, x, p].

    Returns one of the above arrays."""
    f = data[0]
    x = data[1]
    p = data[2]
    q = (-x * np.cos(p / 180 * np.pi) / (x.max() - x.min()) + x.mean()) * mod
    i = 180 / np.pi * x * np.sin(p / 180 * np.pi)

    def outsw(x):
        return {
            3: [f, np.sign(mod) * i, q],
            5: [f, -x, p],
            0: [f, x, p]
        }.get(x, [f, q, i])
    return np.array(outsw(np.abs(mod)))


def test_vna_xp(path, dir, skew, full, mod):
    """Loads a VNA dataset and fits it and its converted counterpart to
    Lorentzians, plotting the result to two separate subplots.

    The fit function is:
    y(x) = b + a * (w / 2)**2 / ((x - f)**2 + (w / 2)**2)

    b is the background offset, a is the peak amplitude, w is the peak width,
    and f is the peak frequency.

    'path' is a string corresponding to the desired datafile;
    'dir' is the data directory;
    'skew' is 1 to subtract a linear background, and 0 to use no slope;
    'full' is 1 to fit the whole dataset, and 0 to only fit between the
    maxima on either side of the peak;
    for details on 'mod' see 'xp_to_quad'."""
    name = glob.glob(os.path.join(dir, path))[0]
    data = np.loadtxt(name, skiprows=1).transpose()
    xdata = xp_to_quad(data, mod)
    f, axarr = plt.subplots(2, sharex=True, figsize=(8, 6))
    ut.plot_lor_fit(data, skew, full, axarr[0])
    ut.plot_lor_fit(xdata, skew, 0, axarr[1])


def load_vna_data(path, dir, mod=0, patt=defNSpatt):
    """Loads all datasets from field sweeps using the VNA corresponding to a
    given string.

    'path' is a string corresponding to the desired datafile;
    'dir' is the data directory;
    'mod' controls whether to convert the amplitude and phase data into
    in-phase and quadrature data, see documentation for
    'xp_to_quad' for full details (default:0, no conversion);
    'patt' is a regular expression pattern to extract temperature and sequence
    information from the filename.

    Returns a list with: a list of datasets; an array of experimental fields;
    an array of experimental temperatures; one of the matching filenames."""
    data, names, archive = ut.get_data_exclude_names(path + '*', dir, 'Fit.txt', 1)

    def matchnum(name):
        [temp, num] = re.match(patt, name).groups()
        return [float(temp), int(num)]
    [temps, nums] = np.array(list(map(matchnum, names))).transpose()

    def fields1(path):
        if archive == 0:
            with open(path) as f:
                first_line = f.readline()
        else:
            with archive.open(path) as f:
                first_line = f.readline().decode('utf-8')
        return float(re.match(filepatt, first_line).groups()[0])
    fields = np.array([fields1(nam) for nam in names])
    sortarr = np.argsort(nums)
    [data, names, fields, temps] = [np.array(data)[sortarr], np.array(names)[
        sortarr], fields[sortarr], temps[sortarr]]
    if mod != 0:
        data = np.array([xp_to_quad(dat, mod) for dat in data])
    if archive != 0:
        archive.close()
    return [data, fields, temps, os.path.split(names[0])[1]]

{'fraction_non_alphanumeric': 0.17390856672158156, 'fraction_numerical': 0.010399505766062602, 'mean_word_length': 2.8836465413834467, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '849328', 'n_tokens_mistral': 3697, 'n_tokens_neox': 3564, 'n_words': 1129}