74 lines
2.3 KiB
Python
74 lines
2.3 KiB
Python
import argparse
|
|
import os
|
|
from tqdm import tqdm
|
|
import tarfile
|
|
|
|
from chargefile import ChargeFile
|
|
from spectrafile import SpectraFile
|
|
from soc_estimation import add_soc_estimate
|
|
from rul import add_rul
|
|
|
|
|
|
if __name__ == "__main__":
|
|
parser = argparse.ArgumentParser("KissExpiramentCreateDataset")
|
|
parser.add_argument('--data', '-d', required=True, help="Data input directory")
|
|
parser.add_argument('--out', '-o', required=True, help="output directory")
|
|
args = parser.parse_args()
|
|
|
|
filenames = [f for f in os.listdir(args.data) if os.path.isfile(os.path.join(args.data, f))]
|
|
charge_filenames = [f for f in filenames if f.startswith("charge") or f.startswith("single_cell_")]
|
|
spectra_filenames = [f for f in filenames if not f.startswith("charge") and not f.startswith("single_cell_") and not f.startswith("voltage_equlaization_") and f != "expiramentlog.csv"]
|
|
|
|
print(f"found {len(spectra_filenames)} spectra")
|
|
print(f"found {len(charge_filenames)} charge/discharge sequences")
|
|
|
|
if not os.path.exists(args.out + ".tmp"):
|
|
os.makedirs(args.out + ".tmp")
|
|
|
|
charge_files = list()
|
|
for filename in charge_filenames:
|
|
charge_files.append(ChargeFile(os.path.join(args.data, filename)))
|
|
charge_files.sort(key=lambda charge_file: charge_file.step)
|
|
|
|
cells = set()
|
|
for filename in tqdm(spectra_filenames, desc="Finding cells"):
|
|
tokens = filename.split('.')[0].split('-')
|
|
if len(tokens) != 3:
|
|
continue
|
|
cellid = int(tokens[1])
|
|
cells.add(cellid)
|
|
|
|
print(f"{len(cells)} cells where involved")
|
|
|
|
spectras = list()
|
|
|
|
for filename in tqdm(spectra_filenames, desc="Resolveing data"):
|
|
tokens = filename.split('.')[0].split('-')
|
|
if len(tokens) != 3:
|
|
continue
|
|
step = int(tokens[0])
|
|
cellid = int(tokens[1])
|
|
substep = int(tokens[2])
|
|
sf = SpectraFile(os.path.join(args.data, filename), cellid, step, substep, charge_files, len(cells))
|
|
spectras.append(sf)
|
|
spectras.sort(key=lambda spectra: spectra.step)
|
|
|
|
add_soc_estimate(spectras)
|
|
add_rul(spectras, cells)
|
|
|
|
for spectra in spectras:
|
|
spectra.write(args.out + ".tmp")
|
|
|
|
try:
|
|
os.remove(f"{args.out}.tar")
|
|
except FileNotFoundError:
|
|
pass
|
|
tar = tarfile.open(f"{args.out}.tar", mode="x")
|
|
for filename in tqdm(os.listdir(args.out + ".tmp"), desc="Saveing data"):
|
|
path = os.path.join(args.out + ".tmp", filename)
|
|
tar.add(path, arcname=os.path.split(path)[-1])
|
|
os.remove(path)
|
|
os.rmdir(args.out + ".tmp")
|
|
tar.close()
|
|
|