From 533b3810655697261cd44f8e3894f93d1b5dc782 Mon Sep 17 00:00:00 2001 From: Tobias Reiter Date: Thu, 17 Jul 2025 13:36:00 +0200 Subject: [PATCH 1/6] Add fit function to tecplot reader --- .../TecPlotReader.py | 123 +++++++++++++++++- .../customParticleDistribution.cpp | 69 +++++++--- include/viennaps/psProcess.hpp | 2 +- 3 files changed, 172 insertions(+), 22 deletions(-) diff --git a/examples/customParticleDistribution/TecPlotReader.py b/examples/customParticleDistribution/TecPlotReader.py index fc669171..5cf3f462 100644 --- a/examples/customParticleDistribution/TecPlotReader.py +++ b/examples/customParticleDistribution/TecPlotReader.py @@ -10,6 +10,7 @@ # Requires: tkinter, matplotlib, numpy, re ##################################################################### +from math import dist import tkinter as tk from tkinter import ttk from tkinter import filedialog, messagebox @@ -18,6 +19,7 @@ import re import argparse import sys +from scipy.optimize import curve_fit # --- Data Handling Functions --- @@ -95,6 +97,7 @@ def trim_data(data, energy, eps=1e-4): def plot_data(data_dict, key, thetaKey, energyKey, eps): d, e = trim_data(data_dict[key], data_dict[energyKey], eps) xx, yy = np.meshgrid(data_dict[thetaKey], e) + plt.figure() plt.pcolormesh(xx, yy, d) plt.colorbar() plt.xlabel(thetaKey) @@ -103,12 +106,91 @@ def plot_data(data_dict, key, thetaKey, energyKey, eps): plt.show() +def line_plot(data_dict, key, thetaKey, energyKey, eps): + d, e = trim_data(data_dict[key], data_dict[energyKey], eps) + theta = data_dict[thetaKey] + fig = plt.figure() + ax = fig.add_subplot(111, projection="3d") + for i in range(len(e)): + ax.plot( + ys=theta, + zs=d[i, :], + xs=e[i], + color=plt.cm.viridis(i / len(e)), + ) + ax.set_ylabel(thetaKey) + ax.set_zlabel(f"Distribution: {key}") + ax.set_xlabel(energyKey) + plt.tight_layout() + plt.show() + + +def fit_data(data_dict, key, thetaKey, energyKey, eps, power_cutoff=10000): + def angle_fit(t, power): + f = np.cos(t) ** power + f = np.abs(np.sin(t) * f) + return f / np.max(f) + + def gaussian(x, mu, sigma): + f = np.exp(-((x - mu) ** 2) / (2 * sigma**2)) + return f / np.max(f) + + def dist_fit(t, e, power, mu, sigma): + return angle_fit(t, power) * gaussian(e, mu, sigma) + + dist, energy = trim_data(data_dict[key], data_dict[energyKey], eps) + theta = data_dict[thetaKey] + theta_fit = np.linspace(np.min(theta), np.max(theta), 1000) + energy_fit = np.linspace(np.min(energy), np.max(energy), 1000) + + powers = np.zeros(len(energy)) + for i in range(len(energy)): + sample = dist[i, :] # / np.max(dist[i, :]) + sample_max = np.max(sample) + popt, _ = curve_fit( + angle_fit, + np.deg2rad(theta), + sample / sample_max, + bounds=([1], [50000]), + p0=[500], + ) + powers[i] = popt[0] + + powers_cut = powers < power_cutoff + popt_p, _ = curve_fit( + lambda x, a, b: a * x + b, + energy[powers_cut], + powers[powers_cut], + bounds=([0, 0], [10, 10000]), + p0=[1, 1], + ) + + dist_t = np.sum(dist, axis=1) + t_max = np.max(dist_t) + popt_g, _ = curve_fit( + gaussian, energy, dist_t / t_max, bounds=([0, 0], [100, 50]), p0=[38, 5] + ) + + xx, yy = np.meshgrid(theta_fit, energy_fit) + d = dist_fit(np.deg2rad(xx), yy, popt_p[1], popt_g[0], popt_g[1]) + plt.figure() + plt.pcolormesh(xx, yy, d, shading="auto") + plt.colorbar(label="Distribution") + plt.xlabel(thetaKey) + plt.ylabel(energyKey) + plt.title( + f"Power {popt_p[1]:.2f}, Gaussian mu {popt_g[0]:.2f}, sigma {popt_g[1]:.2f}" + ) + plt.tight_layout() + plt.show() + + def save_data(data_dict, key, thetaKey, energyKey, eps, outname): d, e = trim_data(data_dict[key], data_dict[energyKey], eps) with open(outname, "w") as f: f.write(f"# {key}\n") - f.write(" ".join(map(str, data_dict[thetaKey])) + "\n") f.write(" ".join(map(str, e)) + "\n") + f.write(" ".join(map(str, data_dict[thetaKey])) + "\n") for row in d: f.write(" ".join(map(str, row)) + "\n") @@ -127,6 +209,7 @@ def __init__(self, parent): self.entry_theta = tk.StringVar(value=self.tk_theta) self.entry_energy = tk.StringVar(value=self.tk_energy) self.entry_zone_lines = tk.IntVar(value=2) + self.line_plot_var = tk.BooleanVar(value=False) # use a built‑in theme style = ttk.Style() @@ -165,6 +248,9 @@ def __init__(self, parent): ttk.Button(frame_buttons, text="Show", command=self.on_plot).pack( side="left", padx=5 ) + ttk.Button(frame_buttons, text="Fit", command=self.on_fit).pack( + side="left", padx=5 + ) ttk.Button(frame_buttons, text="Save...", command=self.on_save).pack( side="left", padx=5 ) @@ -212,7 +298,14 @@ def on_plot(self): return key = self.vars[sel[0]] eps = float(self.entry_eps.get()) - plot_data(self.data, key, self.entry_theta.get(), self.entry_energy.get(), eps) + if self.line_plot_var.get(): + line_plot( + self.data, key, self.entry_theta.get(), self.entry_energy.get(), eps + ) + else: + plot_data( + self.data, key, self.entry_theta.get(), self.entry_energy.get(), eps + ) def on_save(self): sel = self.listbox.curselection() @@ -236,6 +329,16 @@ def on_save(self): ) messagebox.showinfo("Saved", f"Data saved to {out}") + def on_fit(self): + sel = self.listbox.curselection() + if not sel: + messagebox.showwarning("Warning", "No variable selected.") + return + key = self.vars[sel[0]] + eps = float(self.entry_eps.get()) + fit_data(self.data, key, self.entry_theta.get(), self.entry_energy.get(), eps) + # messagebox.showinfo("Fit", f"Fit completed for {key}") + def open_options(self): win = tk.Toplevel(self) win.title("Options") @@ -258,9 +361,25 @@ def open_options(self): entry_zone.pack(side="left", padx=5) frame_zone.pack(pady=10) + # Checkbox for line plot + chk_line_plot = ttk.Checkbutton( + win, + text="Line Plot", + variable=self.line_plot_var, + command=self.toggle_line_plot, + ) + chk_line_plot.pack(anchor="w", padx=15, pady=5) + # Done button to close ttk.Button(win, text="OK", command=win.destroy).pack(pady=10) + def toggle_line_plot(self): + # Toggle the line plot option + if self.line_plot_var.get(): + self.line_plot_var.set(True) + else: + self.line_plot_var.set(False) + def main(): parser = argparse.ArgumentParser( diff --git a/examples/customParticleDistribution/customParticleDistribution.cpp b/examples/customParticleDistribution/customParticleDistribution.cpp index 61b45108..1ead033c 100644 --- a/examples/customParticleDistribution/customParticleDistribution.cpp +++ b/examples/customParticleDistribution/customParticleDistribution.cpp @@ -25,15 +25,7 @@ auto loadDistributionFromFile(const std::string &fileName) { std::string line; // Header std::getline(file, line); - - // Read y-support - std::getline(file, line); - std::istringstream yStream(line); - NumericType yValue; - while (yStream >> yValue) { - distribution.support_y.push_back(yValue); - } - distribution.support_y.shrink_to_fit(); + std::cout << "Reading distribution: " << line << std::endl; // Read x-support std::getline(file, line); @@ -44,6 +36,15 @@ auto loadDistributionFromFile(const std::string &fileName) { } distribution.support_x.shrink_to_fit(); + // Read y-support + std::getline(file, line); + std::istringstream yStream(line); + NumericType yValue; + while (yStream >> yValue) { + distribution.support_y.push_back(yValue); + } + distribution.support_y.shrink_to_fit(); + // Read PDF values size_t rowSize = 0; while (std::getline(file, line)) { @@ -64,20 +65,33 @@ auto loadDistributionFromFile(const std::string &fileName) { distribution.pdf.push_back(pdfRow); } + file.close(); + + // Check if the distribution is valid + if (distribution.pdf.size() != distribution.support_x.size() || + rowSize != distribution.support_y.size()) { + Logger::getInstance() + .addError("Invalid distribution data: " + "Number of rows in PDF does not match x-support size " + "or row size does not match y-support size.") + .print(); + } + return distribution; } template -class BivariateDistributionParticle final - : public viennaray::Particle, +class CustomDistributionParticle final + : public viennaray::Particle, NumericType> { NumericType energy_; + NumericType theta_; Sampling directionNEnergySampling_; // bivariate (2D) sampling instance using // accept-reject sampling public: - explicit BivariateDistributionParticle( + explicit CustomDistributionParticle( const BivariateDistribution &angleEnergyDistribution) { // initialize sampling instances with custom distribution directionNEnergySampling_.setPDF(angleEnergyDistribution.pdf, @@ -91,9 +105,10 @@ class BivariateDistributionParticle final auto sample = directionNEnergySampling_.sample(rngState); assert(sample[1] >= -90. && sample[1] <= 90.); - NumericType theta = constants::degToRad(sample[1]); energy_ = sample[0]; + theta_ = sample[1]; + NumericType theta = constants::degToRad(sample[1]); Vec3D direction{0., 0., 0.}; if constexpr (D == 2) { direction[0] = std::sin(theta); @@ -109,6 +124,21 @@ class BivariateDistributionParticle final return direction; } + void logData(viennaray::DataLog &log) override { + constexpr int num_e = 50; + constexpr int num_theta = 50; + + NumericType t = (theta_ + 15) / 30.; + int theta_bin = static_cast(t * num_theta); + theta_bin = std::clamp(theta_bin, 0, num_theta - 1); + + NumericType e = (energy_ - 10) / (60 - 10); + int e_bin = static_cast(e * num_e); + e_bin = std::clamp(e_bin, 0, num_e - 1); + + log.data[0][theta_bin * num_e + e_bin] += 1; // increment bin + } + void surfaceCollision(NumericType rayWeight, const Vec3D &rayDir, const Vec3D &geomNormal, const unsigned int primID, const int, @@ -188,7 +218,7 @@ class CustomSurfaceModel : public SurfaceModel { int main(int argc, char *argv[]) { - Logger::setLogLevel(LogLevel::INFO); + Logger::setLogLevel(LogLevel::DEBUG); using NumericType = double; constexpr int D = 2; @@ -212,13 +242,14 @@ int main(int argc, char *argv[]) { auto velField = SmartPointer>::New(2); model->setVelocityField(velField); auto distribution = loadDistributionFromFile(filename); - auto particle = - std::make_unique>( - distribution); - model->insertNextParticleType(particle); + auto particle = std::make_unique>( + distribution); + model->insertNextParticleType(particle, 50 * 50); // 50x50 data log size // Run the process - Process(domain, model, 2.0).apply(); + Process p(domain, model, .1); + p.apply(); + p.writeParticleDataLogs("particle_data.txt"); domain->saveVolumeMesh("customParticleDistribution"); diff --git a/include/viennaps/psProcess.hpp b/include/viennaps/psProcess.hpp index e071bcb8..c4ff7b1c 100644 --- a/include/viennaps/psProcess.hpp +++ b/include/viennaps/psProcess.hpp @@ -56,7 +56,7 @@ class Process final : public ProcessBase { for (std::size_t i = 0; i < particleDataLogs.size(); i++) { if (!particleDataLogs[i].data.empty()) { - file << "particle" << i << "_data "; + file << "particle" << i << "_data\n"; for (std::size_t j = 0; j < particleDataLogs[i].data[0].size(); j++) { file << particleDataLogs[i].data[0][j] << " "; } From 9762d21ee3c1540eb9313e2d0b6e4dc4cb84f943 Mon Sep 17 00:00:00 2001 From: Tobias Reiter Date: Fri, 18 Jul 2025 10:28:44 +0200 Subject: [PATCH 2/6] Improve TecPlotReader fit function --- .../TecPlotReader.py | 307 +++++++++++++----- 1 file changed, 234 insertions(+), 73 deletions(-) diff --git a/examples/customParticleDistribution/TecPlotReader.py b/examples/customParticleDistribution/TecPlotReader.py index 5cf3f462..f3896fe5 100644 --- a/examples/customParticleDistribution/TecPlotReader.py +++ b/examples/customParticleDistribution/TecPlotReader.py @@ -10,12 +10,12 @@ # Requires: tkinter, matplotlib, numpy, re ##################################################################### -from math import dist import tkinter as tk from tkinter import ttk from tkinter import filedialog, messagebox import numpy as np import matplotlib.pyplot as plt +from matplotlib.widgets import TextBox, Button import re import argparse import sys @@ -57,7 +57,7 @@ def read_tecplot_data(filename, thetaKey, energyKey, zone_lines_max=2): I, J = int(mI.group(1)), int(mJ.group(1)) zone_lines_read += 1 if zone_lines_read == zone_lines_max: - mode = 2 + mode = 2 # switch to data mode elif mode == 2: var = variables[var_counter] parts = line.split() @@ -76,6 +76,12 @@ def read_tecplot_data(filename, thetaKey, energyKey, zone_lines_max=2): # convert to arrays for k in data_dict: data_dict[k] = np.array(data_dict[k]) + + if thetaKey not in data_dict or energyKey not in data_dict: + raise ValueError( + f"Required keys '{thetaKey}' or '{energyKey}' not found in the data." + ) + # extract 1D theta and energy data_dict[thetaKey] = data_dict[thetaKey][0, :] data_dict[energyKey] = data_dict[energyKey][:, 0] @@ -99,7 +105,7 @@ def plot_data(data_dict, key, thetaKey, energyKey, eps): xx, yy = np.meshgrid(data_dict[thetaKey], e) plt.figure() plt.pcolormesh(xx, yy, d) - plt.colorbar() + plt.colorbar(label="Distribution") plt.xlabel(thetaKey) plt.ylabel(energyKey) plt.title(f"Distribution: {key}") @@ -125,10 +131,20 @@ def line_plot(data_dict, key, thetaKey, energyKey, eps): plt.show() -def fit_data(data_dict, key, thetaKey, energyKey, eps, power_cutoff=10000): +def fit_data( + data_dict, + key, + thetaKey, + energyKey, + eps, + power_cutoff=10000, + useSin=False, + verbose=True, +): def angle_fit(t, power): f = np.cos(t) ** power - f = np.abs(np.sin(t) * f) + if useSin: + f = np.abs(np.sin(t) * f) return f / np.max(f) def gaussian(x, mu, sigma): @@ -147,14 +163,75 @@ def dist_fit(t, e, power, mu, sigma): for i in range(len(energy)): sample = dist[i, :] # / np.max(dist[i, :]) sample_max = np.max(sample) - popt, _ = curve_fit( - angle_fit, - np.deg2rad(theta), - sample / sample_max, - bounds=([1], [50000]), - p0=[500], + if sample_max > 0: + sample = sample / sample_max + popt, _ = curve_fit( + angle_fit, + np.deg2rad(theta), + sample, + bounds=([1], [50000]), + p0=[500], + ) + powers[i] = popt[0] + else: + powers[i] = 50000 + + if verbose: + fig, ax = plt.subplots() + plt.subplots_adjust(bottom=0.3) # leave room at bottom + [line_dist] = ax.plot(theta, dist[0, :], label="Distribution") + [line_fit] = ax.plot( + theta_fit, + angle_fit(np.deg2rad(theta_fit), powers[0]) * np.max(dist[0, :]), + label="Fit", ) - powers[i] = popt[0] + ax.set_xlabel(thetaKey) + ax.set_title(f"Energy = {energy[0]:.2f} eV, Power = {powers[0]:.2f}") + + # TextBox for the integer, centered + axbox = plt.axes([0.4, 0.15, 0.2, 0.05]) + text_box = TextBox(axbox, "", initial=str(0)) + + # “–” button + axdec = plt.axes([0.33, 0.15, 0.05, 0.05]) + btn_dec = Button(axdec, "–") + + # “+” button + axinc = plt.axes([0.62, 0.15, 0.05, 0.05]) + btn_inc = Button(axinc, "+") + + def update_plot(): + try: + f = int(text_box.text) + except ValueError: + return + line_dist.set_ydata(dist[f, :]) + line_fit.set_ydata( + angle_fit(np.deg2rad(theta_fit), powers[f]) * np.max(dist[f, :]) + ) + ax.set_title(f"Energy = {energy[f]:.2f} eV, Power = {powers[f]:.2f}") + ax.relim() + ax.autoscale_view() + fig.canvas.draw_idle() + + def on_inc(event): + v = int(text_box.text) + 1 + if v >= len(energy): + v = len(energy) - 1 + text_box.set_val(str(v)) + update_plot() + + def on_dec(event): + v = int(text_box.text) - 1 + if v < 0: + v = 0 + text_box.set_val(str(v)) + update_plot() + + # wire up callbacks + btn_inc.on_clicked(on_inc) + btn_dec.on_clicked(on_dec) + text_box.on_submit(lambda text: update_plot()) powers_cut = powers < power_cutoff popt_p, _ = curve_fit( @@ -165,11 +242,45 @@ def dist_fit(t, e, power, mu, sigma): p0=[1, 1], ) + if verbose: + plt.figure() + plt.plot(energy, powers, "-x", label="Powers") + plt.plot( + energy[powers_cut], + popt_p[0] * energy[powers_cut] + popt_p[1], + "o--", + label="Power Fit", + ) + plt.axhline(power_cutoff, color="red", linestyle="--", label="Cutoff") + plt.xlabel(energyKey) + plt.ylabel("Power") + plt.title(f"Power Fit: a={popt_p[0]:.2f}, b={popt_p[1]:.2f}") + plt.legend() + dist_t = np.sum(dist, axis=1) t_max = np.max(dist_t) - popt_g, _ = curve_fit( - gaussian, energy, dist_t / t_max, bounds=([0, 0], [100, 50]), p0=[38, 5] - ) + if t_max > 0: + mu_0 = np.trapezoid(dist_t * energy, x=energy) + popt_g, _ = curve_fit( + gaussian, energy, dist_t / t_max, bounds=([0, 0], [100, 50]), p0=[mu_0, 5] + ) + else: + return -1 + + if verbose: + fig, ax_e = plt.subplots() + ax_e.plot(energy, dist_t, label="Cumulative Distribution") + ax_e.plot( + energy_fit, + gaussian(energy_fit, popt_g[0], popt_g[1]) * t_max, + label="Gaussian Fit", + ) + ax_e.set_xlabel(energyKey) + ax_e.set_ylabel("Cumulative Distribution") + ax_e.set_title(f"Gaussian Fit: mu={popt_g[0]:.2f}, sigma={popt_g[1]:.2f}") + ax_e.legend() + + print(f"Power {popt_p[1]:.2f}, Gaussian mu {popt_g[0]:.2f}, sigma {popt_g[1]:.2f}") xx, yy = np.meshgrid(theta_fit, energy_fit) d = dist_fit(np.deg2rad(xx), yy, popt_p[1], popt_g[0], popt_g[1]) @@ -184,6 +295,8 @@ def dist_fit(t, e, power, mu, sigma): plt.tight_layout() plt.show() + return 0 + def save_data(data_dict, key, thetaKey, energyKey, eps, outname): d, e = trim_data(data_dict[key], data_dict[energyKey], eps) @@ -200,77 +313,83 @@ class TecplotGUI(ttk.Frame): def __init__(self, parent): super().__init__(parent, padding=10) parent.title("Tecplot Distribution Viewer") - parent.minsize(500, 400) + parent.minsize(400, 300) self.data = {} self.vars = [] + # use a built‑in theme + # style = ttk.Style() + # style.theme_use("clam") + + # Options Input self.tk_theta = "THETA (DEG)" self.tk_energy = "ENERGY (EV)" self.entry_theta = tk.StringVar(value=self.tk_theta) self.entry_energy = tk.StringVar(value=self.tk_energy) self.entry_zone_lines = tk.IntVar(value=2) self.line_plot_var = tk.BooleanVar(value=False) - - # use a built‑in theme - style = ttk.Style() - style.theme_use("clam") + self.fit_power_cutoff = tk.IntVar(value=10000) + self.fit_use_sin = tk.BooleanVar(value=False) + self.fit_verbose = tk.BooleanVar(value=False) + self.path_fn = "" # File input - ttk.Label(self, text="Filename:").pack(anchor="w") frame_file = ttk.Frame(self) - self.entry_file = ttk.Entry(frame_file, width=40) + ttk.Label(frame_file, text="Filename:").pack(anchor="w") + self.entry_file = ttk.Entry(frame_file, width=20) self.entry_file.pack(side="left", padx=5) ttk.Button(frame_file, text="Browse", command=self.browse_file).pack( side="left", padx=5 ) - ttk.Button(frame_file, text="\u21bb", command=self.load_data, width=4).pack( - side="left", padx=5 - ) frame_file.pack(pady=5) # Variable list - ttk.Label(self, text="Select Variable:").pack(anchor="w") + frame_main = ttk.Frame(self) + ttk.Label(frame_main, text="Select Variable:").grid(row=0, column=0) self.listbox = tk.Listbox( - self, background="white", selectmode=tk.SINGLE, foreground="black" + frame_main, background="gray", selectmode=tk.SINGLE, foreground="black" ) - self.listbox.pack(fill="both", expand=True, padx=5) + self.listbox.grid(row=1, column=0, padx=5, rowspan=6) # Trim eps - frame_eps = ttk.Frame(self) - ttk.Label(frame_eps, text="Trim eps:").pack(side="left") - self.entry_eps = ttk.Entry(frame_eps, width=10) + ttk.Label(frame_main, text="Trim eps:").grid(row=1, column=1, padx=5) + self.entry_eps = ttk.Entry(frame_main, width=10) self.entry_eps.insert(0, "1e-4") - self.entry_eps.pack(side="left", padx=5) - frame_eps.pack(pady=5) + self.entry_eps.grid(row=2, column=1, padx=5) # Buttons - frame_buttons = ttk.Frame(self) - ttk.Button(frame_buttons, text="Show", command=self.on_plot).pack( - side="left", padx=5 + ttk.Button(frame_main, text="Show", command=self.on_plot).grid( + row=3, column=1, padx=5 ) - ttk.Button(frame_buttons, text="Fit", command=self.on_fit).pack( - side="left", padx=5 + ttk.Button(frame_main, text="Fit", command=self.on_fit).grid( + row=4, column=1, padx=5 ) - ttk.Button(frame_buttons, text="Save...", command=self.on_save).pack( - side="left", padx=5 + ttk.Button(frame_main, text="Save", command=self.on_save).grid( + row=5, column=1, padx=5 ) - frame_buttons.pack(pady=10) # Extended options button - btn_ext = ttk.Button(frame_buttons, text="Options", command=self.open_options) - btn_ext.pack(side="left", padx=5) + btn_ext = ttk.Button( + frame_main, text="Options", command=self.open_options + ).grid(row=6, column=1, padx=5) + frame_main.pack(pady=5) def browse_file(self): f = filedialog.askopenfilename( - filetypes=[("TecPlot files", "*.pdt"), ("All files", "*")] + filetypes=[("TecPlot files", "*.pdt"), ("All files", "*")], + title="Select TecPlot File", + initialdir=".", ) if f: self.entry_file.delete(0, tk.END) + # insert the filename into the entry + self.path_fn = f + f = f.split("/")[-1] if "/" in f else f.split("\\")[-1] self.entry_file.insert(0, f) self.load_data() def load_data(self): - fn = self.entry_file.get().strip() + fn = self.path_fn.strip() self.tk_theta = self.entry_theta.get().strip() self.tk_energy = self.entry_energy.get().strip() zone_lines = self.entry_zone_lines.get() @@ -336,49 +455,74 @@ def on_fit(self): return key = self.vars[sel[0]] eps = float(self.entry_eps.get()) - fit_data(self.data, key, self.entry_theta.get(), self.entry_energy.get(), eps) - # messagebox.showinfo("Fit", f"Fit completed for {key}") + power_cutoff = self.fit_power_cutoff.get() + use_sine = self.fit_use_sin.get() + verbose = self.fit_verbose.get() + res = fit_data( + self.data, + key, + self.entry_theta.get(), + self.entry_energy.get(), + eps, + power_cutoff=power_cutoff, + useSin=use_sine, + verbose=verbose, + ) + if res == -1: + messagebox.showerror("Error", "Could not fit data.") def open_options(self): win = tk.Toplevel(self) win.title("Options") - win.geometry("300x200") + win.geometry("300x400") + win.resizable(False, False) # Keys - frame_keys = tk.Frame(win) - tk.Label(frame_keys, text="Theta Key:").grid(row=0, column=0) - entry_t = tk.Entry(frame_keys, textvariable=self.entry_theta) - entry_t.grid(row=0, column=1) - tk.Label(frame_keys, text="Energy Key:").grid(row=1, column=0) - entry_e = tk.Entry(frame_keys, textvariable=self.entry_energy) - entry_e.grid(row=1, column=1) + frame_keys = ttk.Frame(win) + ttk.Label(frame_keys, text="Parsing:").grid(row=0, columnspan=2, pady=5) + ttk.Label(frame_keys, text="Theta Key:").grid(row=1, column=0) + entry_t = ttk.Entry(frame_keys, textvariable=self.entry_theta, width=20) + entry_t.grid(row=1, column=1) + ttk.Label(frame_keys, text="Energy Key:").grid(row=2, column=0) + entry_e = ttk.Entry(frame_keys, textvariable=self.entry_energy, width=20) + entry_e.grid(row=2, column=1) + ttk.Label(frame_keys, text="Zone Lines:").grid(row=3, column=0) + entry_zone = ttk.Entry(frame_keys, textvariable=self.entry_zone_lines, width=20) + entry_zone.grid(row=3, column=1) frame_keys.pack(pady=10) - # Zone lines - frame_zone = tk.Frame(win) - tk.Label(frame_zone, text="Zone Lines:").pack(side="left") - entry_zone = tk.Entry(frame_zone, textvariable=self.entry_zone_lines, width=5) - entry_zone.pack(side="left", padx=5) - frame_zone.pack(pady=10) - # Checkbox for line plot + frame_plot = ttk.Frame(win) + ttk.Label(frame_plot, text="Plot Options:").pack(anchor="center") chk_line_plot = ttk.Checkbutton( - win, + frame_plot, text="Line Plot", variable=self.line_plot_var, - command=self.toggle_line_plot, ) - chk_line_plot.pack(anchor="w", padx=15, pady=5) + chk_line_plot.pack(anchor="center", pady=5) + frame_plot.pack(pady=10) + + frame_fit = ttk.Frame(win) + ttk.Label(frame_fit, text="Fit Options:").grid(row=0, columnspan=2) + ttk.Label(frame_fit, text="Power Cutoff:").grid(row=1, column=0) + entry_power = ttk.Entry(frame_fit, textvariable=self.fit_power_cutoff, width=20) + entry_power.grid(row=1, column=1) + chk_fit_sin = ttk.Checkbutton( + frame_fit, + text="Use Sine in Fit", + variable=self.fit_use_sin, + ) + chk_fit_sin.grid(row=2, columnspan=2, pady=5) + chk_fit_verbose = ttk.Checkbutton( + frame_fit, + text="Show All Plots", + variable=self.fit_verbose, + ) + chk_fit_verbose.grid(row=3, columnspan=2, pady=5) + frame_fit.pack(pady=10) # Done button to close - ttk.Button(win, text="OK", command=win.destroy).pack(pady=10) - - def toggle_line_plot(self): - # Toggle the line plot option - if self.line_plot_var.get(): - self.line_plot_var.set(True) - else: - self.line_plot_var.set(False) + ttk.Button(win, text="OK", command=win.destroy).pack(anchor="s", pady=5) def main(): @@ -393,6 +537,7 @@ def main(): "-o", "--output", help="output filename", metavar="OUT", default="output.txt" ) parser.add_argument("--eps", help="trim epsilon", type=float, default=1e-4) + parser.add_argument("--fit", help="Fit distribution to data", action="store_true") args = parser.parse_args() @@ -415,6 +560,22 @@ def main(): energyKey=energy_key, ) + if args.fit: + if args.input == "ALL": + print("Error: Variable input need for fit.") + return + + if args.input not in variables: + print( + f"Error: variable '{args.input}' not found in {args.file}", + ) + return + + res = fit_data(data_dict, args.input, theta_key, energy_key, args.eps) + if res == -1: + print("Error fitting data.") + return + outFileName = args.output if outFileName.endswith(".txt"): outFileName = outFileName[:-4] From 16d4e866b8299d9ac1a5ed673a33cbf8c6675a22 Mon Sep 17 00:00:00 2001 From: Tobias Reiter Date: Fri, 18 Jul 2025 10:57:17 +0200 Subject: [PATCH 3/6] Small fixes --- .../TecPlotReader.py | 43 +++++++++---------- 1 file changed, 20 insertions(+), 23 deletions(-) diff --git a/examples/customParticleDistribution/TecPlotReader.py b/examples/customParticleDistribution/TecPlotReader.py index f3896fe5..c439f02b 100644 --- a/examples/customParticleDistribution/TecPlotReader.py +++ b/examples/customParticleDistribution/TecPlotReader.py @@ -6,8 +6,9 @@ # Usage: # python TecPlotReader.py -f data.pdt # python TecPlotReader.py -f data.pdt -i variable_name -o output.txt +# python TecPlotReader.py -f data.pdt -i variable_name --fit # python TecPlotReader.py # to open the GUI -# Requires: tkinter, matplotlib, numpy, re +# Requires: tkinter, matplotlib, numpy, scipy ##################################################################### import tkinter as tk @@ -161,20 +162,19 @@ def dist_fit(t, e, power, mu, sigma): powers = np.zeros(len(energy)) for i in range(len(energy)): - sample = dist[i, :] # / np.max(dist[i, :]) + sample = dist[i, :] sample_max = np.max(sample) if sample_max > 0: - sample = sample / sample_max popt, _ = curve_fit( angle_fit, np.deg2rad(theta), - sample, - bounds=([1], [50000]), + sample / sample_max, + bounds=([1], [power_cutoff]), p0=[500], ) powers[i] = popt[0] else: - powers[i] = 50000 + powers[i] = power_cutoff if verbose: fig, ax = plt.subplots() @@ -203,6 +203,8 @@ def dist_fit(t, e, power, mu, sigma): def update_plot(): try: f = int(text_box.text) + if f < 0 or f >= len(energy): + raise ValueError("Index out of bounds") except ValueError: return line_dist.set_ydata(dist[f, :]) @@ -215,16 +217,12 @@ def update_plot(): fig.canvas.draw_idle() def on_inc(event): - v = int(text_box.text) + 1 - if v >= len(energy): - v = len(energy) - 1 + v = min(int(text_box.text) + 1, len(energy) - 1) text_box.set_val(str(v)) update_plot() def on_dec(event): - v = int(text_box.text) - 1 - if v < 0: - v = 0 + v = max(int(text_box.text) - 1, 0) text_box.set_val(str(v)) update_plot() @@ -268,24 +266,23 @@ def on_dec(event): return -1 if verbose: - fig, ax_e = plt.subplots() - ax_e.plot(energy, dist_t, label="Cumulative Distribution") - ax_e.plot( + plt.plot(energy, dist_t, label="Cumulative Distribution") + plt.plot( energy_fit, - gaussian(energy_fit, popt_g[0], popt_g[1]) * t_max, + gaussian(energy_fit, *popt_g) * t_max, label="Gaussian Fit", ) - ax_e.set_xlabel(energyKey) - ax_e.set_ylabel("Cumulative Distribution") - ax_e.set_title(f"Gaussian Fit: mu={popt_g[0]:.2f}, sigma={popt_g[1]:.2f}") - ax_e.legend() + plt.xlabel(energyKey) + plt.ylabel("Cumulative Distribution") + plt.title(f"Gaussian Fit: mu={popt_g[0]:.2f}, sigma={popt_g[1]:.2f}") + plt.legend() print(f"Power {popt_p[1]:.2f}, Gaussian mu {popt_g[0]:.2f}, sigma {popt_g[1]:.2f}") - xx, yy = np.meshgrid(theta_fit, energy_fit) - d = dist_fit(np.deg2rad(xx), yy, popt_p[1], popt_g[0], popt_g[1]) + tt, ee = np.meshgrid(theta_fit, energy_fit) + d = dist_fit(np.deg2rad(tt), ee, popt_p[1], popt_g[0], popt_g[1]) plt.figure() - plt.pcolormesh(xx, yy, d, shading="auto") + plt.pcolormesh(tt, ee, d, shading="auto") plt.colorbar(label="Distribution") plt.xlabel(thetaKey) plt.ylabel(energyKey) From 4c6d4a47b98bdace3c86c4a077afe58bbde9c677 Mon Sep 17 00:00:00 2001 From: Tobias Reiter Date: Fri, 18 Jul 2025 11:22:49 +0200 Subject: [PATCH 4/6] Move TecPlotReader to python dir --- {examples/customParticleDistribution => python}/TecPlotReader.py | 0 1 file changed, 0 insertions(+), 0 deletions(-) rename {examples/customParticleDistribution => python}/TecPlotReader.py (100%) diff --git a/examples/customParticleDistribution/TecPlotReader.py b/python/TecPlotReader.py similarity index 100% rename from examples/customParticleDistribution/TecPlotReader.py rename to python/TecPlotReader.py From fe449062708196ae1b27307fa5c43e40448de881 Mon Sep 17 00:00:00 2001 From: Tobias Reiter Date: Fri, 18 Jul 2025 12:48:20 +0200 Subject: [PATCH 5/6] Add von Mises function for fit --- python/TecPlotReader.py | 115 ++++++++++++++++++++++++++-------------- 1 file changed, 76 insertions(+), 39 deletions(-) diff --git a/python/TecPlotReader.py b/python/TecPlotReader.py index c439f02b..80874713 100644 --- a/python/TecPlotReader.py +++ b/python/TecPlotReader.py @@ -21,6 +21,7 @@ import argparse import sys from scipy.optimize import curve_fit +from scipy import stats # --- Data Handling Functions --- @@ -138,29 +139,40 @@ def fit_data( thetaKey, energyKey, eps, - power_cutoff=10000, + angle_p_cutoff=10000, useSin=False, - verbose=True, + verbose=False, + vm=False, ): - def angle_fit(t, power): + def _powerCos(t, power): f = np.cos(t) ** power if useSin: f = np.abs(np.sin(t) * f) return f / np.max(f) - def gaussian(x, mu, sigma): + def _vonMises(t, kappa): + f = stats.vonmises.pdf(t, kappa=kappa, loc=0, scale=np.pi / 2) + if useSin: + f = np.abs(np.sin(t) * f) + return f / np.max(f) + + def _gaussian(x, mu, sigma): f = np.exp(-((x - mu) ** 2) / (2 * sigma**2)) return f / np.max(f) - def dist_fit(t, e, power, mu, sigma): - return angle_fit(t, power) * gaussian(e, mu, sigma) + angle_fit = _vonMises if vm else _powerCos + angle_param = "Kappa" if vm else "Power" + + def dist_fit(t, e, p, mu, sigma): + return angle_fit(t, p) * _gaussian(e, mu, sigma) dist, energy = trim_data(data_dict[key], data_dict[energyKey], eps) theta = data_dict[thetaKey] theta_fit = np.linspace(np.min(theta), np.max(theta), 1000) energy_fit = np.linspace(np.min(energy), np.max(energy), 1000) - powers = np.zeros(len(energy)) + # fit angle distribution + angle_p = np.zeros(len(energy)) for i in range(len(energy)): sample = dist[i, :] sample_max = np.max(sample) @@ -169,12 +181,12 @@ def dist_fit(t, e, power, mu, sigma): angle_fit, np.deg2rad(theta), sample / sample_max, - bounds=([1], [power_cutoff]), - p0=[500], + bounds=([1], [angle_p_cutoff]), + p0=[100], ) - powers[i] = popt[0] + angle_p[i] = popt[0] else: - powers[i] = power_cutoff + angle_p[i] = angle_p_cutoff if verbose: fig, ax = plt.subplots() @@ -182,11 +194,11 @@ def dist_fit(t, e, power, mu, sigma): [line_dist] = ax.plot(theta, dist[0, :], label="Distribution") [line_fit] = ax.plot( theta_fit, - angle_fit(np.deg2rad(theta_fit), powers[0]) * np.max(dist[0, :]), + angle_fit(np.deg2rad(theta_fit), angle_p[0]) * np.max(dist[0, :]), label="Fit", ) ax.set_xlabel(thetaKey) - ax.set_title(f"Energy = {energy[0]:.2f} eV, Power = {powers[0]:.2f}") + ax.set_title(f"Energy = {energy[0]:.2f} eV, {angle_param} = {angle_p[0]:.2f}") # TextBox for the integer, centered axbox = plt.axes([0.4, 0.15, 0.2, 0.05]) @@ -194,7 +206,7 @@ def dist_fit(t, e, power, mu, sigma): # “–” button axdec = plt.axes([0.33, 0.15, 0.05, 0.05]) - btn_dec = Button(axdec, "–") + btn_dec = Button(axdec, "-") # “+” button axinc = plt.axes([0.62, 0.15, 0.05, 0.05]) @@ -209,9 +221,11 @@ def update_plot(): return line_dist.set_ydata(dist[f, :]) line_fit.set_ydata( - angle_fit(np.deg2rad(theta_fit), powers[f]) * np.max(dist[f, :]) + angle_fit(np.deg2rad(theta_fit), angle_p[f]) * np.max(dist[f, :]) + ) + ax.set_title( + f"Energy = {energy[f]:.2f} eV, {angle_param} = {angle_p[f]:.2f}" ) - ax.set_title(f"Energy = {energy[f]:.2f} eV, Power = {powers[f]:.2f}") ax.relim() ax.autoscale_view() fig.canvas.draw_idle() @@ -231,45 +245,48 @@ def on_dec(event): btn_dec.on_clicked(on_dec) text_box.on_submit(lambda text: update_plot()) - powers_cut = powers < power_cutoff + # fit resulting parameter distribution + angle_p_cut = angle_p < (angle_p_cutoff - 1) popt_p, _ = curve_fit( lambda x, a, b: a * x + b, - energy[powers_cut], - powers[powers_cut], + energy[angle_p_cut], + angle_p[angle_p_cut], bounds=([0, 0], [10, 10000]), p0=[1, 1], ) if verbose: plt.figure() - plt.plot(energy, powers, "-x", label="Powers") + plt.plot(energy, angle_p, "-x", label=angle_param) plt.plot( - energy[powers_cut], - popt_p[0] * energy[powers_cut] + popt_p[1], + energy[angle_p_cut], + popt_p[0] * energy[angle_p_cut] + popt_p[1], "o--", - label="Power Fit", + label=f"{angle_param} Fit", ) - plt.axhline(power_cutoff, color="red", linestyle="--", label="Cutoff") + plt.axhline(angle_p_cutoff, color="red", linestyle="--", label="Cutoff") plt.xlabel(energyKey) - plt.ylabel("Power") - plt.title(f"Power Fit: a={popt_p[0]:.2f}, b={popt_p[1]:.2f}") + plt.ylabel(angle_param) + plt.title(f"{angle_param} Fit: a={popt_p[0]:.2f}, b={popt_p[1]:.2f}") plt.legend() + # fit energy distribution dist_t = np.sum(dist, axis=1) t_max = np.max(dist_t) if t_max > 0: mu_0 = np.trapezoid(dist_t * energy, x=energy) popt_g, _ = curve_fit( - gaussian, energy, dist_t / t_max, bounds=([0, 0], [100, 50]), p0=[mu_0, 5] + _gaussian, energy, dist_t / t_max, bounds=([0, 0], [100, 50]), p0=[mu_0, 5] ) else: return -1 if verbose: + plt.figure() plt.plot(energy, dist_t, label="Cumulative Distribution") plt.plot( energy_fit, - gaussian(energy_fit, *popt_g) * t_max, + _gaussian(energy_fit, *popt_g) * t_max, label="Gaussian Fit", ) plt.xlabel(energyKey) @@ -277,7 +294,9 @@ def on_dec(event): plt.title(f"Gaussian Fit: mu={popt_g[0]:.2f}, sigma={popt_g[1]:.2f}") plt.legend() - print(f"Power {popt_p[1]:.2f}, Gaussian mu {popt_g[0]:.2f}, sigma {popt_g[1]:.2f}") + print( + f"{angle_param} {popt_p[1]:.2f}, Gaussian mu {popt_g[0]:.2f}, sigma {popt_g[1]:.2f}" + ) tt, ee = np.meshgrid(theta_fit, energy_fit) d = dist_fit(np.deg2rad(tt), ee, popt_p[1], popt_g[0], popt_g[1]) @@ -287,7 +306,7 @@ def on_dec(event): plt.xlabel(thetaKey) plt.ylabel(energyKey) plt.title( - f"Power {popt_p[1]:.2f}, Gaussian mu {popt_g[0]:.2f}, sigma {popt_g[1]:.2f}" + f"{angle_param} {popt_p[1]:.2f}, Gaussian mu {popt_g[0]:.2f}, sigma {popt_g[1]:.2f}" ) plt.tight_layout() plt.show() @@ -325,9 +344,10 @@ def __init__(self, parent): self.entry_energy = tk.StringVar(value=self.tk_energy) self.entry_zone_lines = tk.IntVar(value=2) self.line_plot_var = tk.BooleanVar(value=False) - self.fit_power_cutoff = tk.IntVar(value=10000) + self.fit_angle_p_cutoff = tk.IntVar(value=10000) self.fit_use_sin = tk.BooleanVar(value=False) self.fit_verbose = tk.BooleanVar(value=False) + self.fit_function = tk.StringVar(value="powerCosine") self.path_fn = "" # File input @@ -452,21 +472,30 @@ def on_fit(self): return key = self.vars[sel[0]] eps = float(self.entry_eps.get()) - power_cutoff = self.fit_power_cutoff.get() + angle_p_cutoff = self.fit_angle_p_cutoff.get() use_sine = self.fit_use_sin.get() verbose = self.fit_verbose.get() + fit_function = self.fit_function.get() + if fit_function == "vonMises": + vm = True + elif fit_function == "powerCosine": + vm = False + else: + messagebox.showerror("Error", "Unknown fit function selected.") + return res = fit_data( self.data, key, self.entry_theta.get(), self.entry_energy.get(), eps, - power_cutoff=power_cutoff, + angle_p_cutoff=angle_p_cutoff, useSin=use_sine, verbose=verbose, + vm=vm, ) if res == -1: - messagebox.showerror("Error", "Could not fit data.") + messagebox.showerror("Error", "Failed to fit data.") def open_options(self): win = tk.Toplevel(self) @@ -501,21 +530,29 @@ def open_options(self): frame_fit = ttk.Frame(win) ttk.Label(frame_fit, text="Fit Options:").grid(row=0, columnspan=2) - ttk.Label(frame_fit, text="Power Cutoff:").grid(row=1, column=0) - entry_power = ttk.Entry(frame_fit, textvariable=self.fit_power_cutoff, width=20) - entry_power.grid(row=1, column=1) + ttk.Combobox( + frame_fit, + textvariable=self.fit_function, + values=["powerCosine", "vonMises"], + state="readonly", + ).grid(row=1, columnspan=2, pady=5) + ttk.Label(frame_fit, text="Cutoff:").grid(row=2, column=0) + entry_power = ttk.Entry( + frame_fit, textvariable=self.fit_angle_p_cutoff, width=20 + ) + entry_power.grid(row=2, column=1) chk_fit_sin = ttk.Checkbutton( frame_fit, text="Use Sine in Fit", variable=self.fit_use_sin, ) - chk_fit_sin.grid(row=2, columnspan=2, pady=5) + chk_fit_sin.grid(row=3, columnspan=2, pady=5) chk_fit_verbose = ttk.Checkbutton( frame_fit, text="Show All Plots", variable=self.fit_verbose, ) - chk_fit_verbose.grid(row=3, columnspan=2, pady=5) + chk_fit_verbose.grid(row=4, columnspan=2, pady=5) frame_fit.pack(pady=10) # Done button to close From 55bc450feae4c5855429744c97f306b5579d5865 Mon Sep 17 00:00:00 2001 From: Tobias Reiter Date: Fri, 18 Jul 2025 13:59:49 +0200 Subject: [PATCH 6/6] Add fit weighting parameters --- python/TecPlotReader.py | 22 ++++++++++++++++++---- 1 file changed, 18 insertions(+), 4 deletions(-) diff --git a/python/TecPlotReader.py b/python/TecPlotReader.py index 80874713..ea9ce49b 100644 --- a/python/TecPlotReader.py +++ b/python/TecPlotReader.py @@ -143,6 +143,8 @@ def fit_data( useSin=False, verbose=False, vm=False, + alpha=0.0, + beta=2.0, ): def _powerCos(t, power): f = np.cos(t) ** power @@ -174,7 +176,9 @@ def dist_fit(t, e, p, mu, sigma): # fit angle distribution angle_p = np.zeros(len(energy)) for i in range(len(energy)): - sample = dist[i, :] + sample = dist[i, :] * ( + beta - (beta - 1) * np.exp(-alpha * np.abs(np.deg2rad(theta))) + ) sample_max = np.max(sample) if sample_max > 0: popt, _ = curve_fit( @@ -348,6 +352,8 @@ def __init__(self, parent): self.fit_use_sin = tk.BooleanVar(value=False) self.fit_verbose = tk.BooleanVar(value=False) self.fit_function = tk.StringVar(value="powerCosine") + self.fit_alpha = tk.DoubleVar(value=0.0) + self.fit_beta = tk.DoubleVar(value=2.0) self.path_fn = "" # File input @@ -493,6 +499,8 @@ def on_fit(self): useSin=use_sine, verbose=verbose, vm=vm, + alpha=self.fit_alpha.get(), + beta=self.fit_beta.get(), ) if res == -1: messagebox.showerror("Error", "Failed to fit data.") @@ -500,7 +508,7 @@ def on_fit(self): def open_options(self): win = tk.Toplevel(self) win.title("Options") - win.geometry("300x400") + win.geometry("300x470") win.resizable(False, False) # Keys @@ -541,18 +549,24 @@ def open_options(self): frame_fit, textvariable=self.fit_angle_p_cutoff, width=20 ) entry_power.grid(row=2, column=1) + ttk.Label(frame_fit, text="Alpha:").grid(row=3, column=0) + entry_alpha = ttk.Entry(frame_fit, textvariable=self.fit_alpha, width=20) + entry_alpha.grid(row=3, column=1) + ttk.Label(frame_fit, text="Beta:").grid(row=4, column=0) + entry_beta = ttk.Entry(frame_fit, textvariable=self.fit_beta, width=20) + entry_beta.grid(row=4, column=1) chk_fit_sin = ttk.Checkbutton( frame_fit, text="Use Sine in Fit", variable=self.fit_use_sin, ) - chk_fit_sin.grid(row=3, columnspan=2, pady=5) + chk_fit_sin.grid(row=5, columnspan=2, pady=5) chk_fit_verbose = ttk.Checkbutton( frame_fit, text="Show All Plots", variable=self.fit_verbose, ) - chk_fit_verbose.grid(row=4, columnspan=2, pady=5) + chk_fit_verbose.grid(row=6, columnspan=2, pady=5) frame_fit.pack(pady=10) # Done button to close