Initial commit for project 3..... A lot happend
This commit is contained in:
Executable
+87
@@ -0,0 +1,87 @@
|
||||
#!/bin/bash
|
||||
|
||||
PUEUE=true
|
||||
NORMAL_MANY_PARTICLE_SIMULATION=true
|
||||
OSCILLATING_FIELD_SIMULATION=false
|
||||
|
||||
echo "Removing old results..."
|
||||
mkdir -p results
|
||||
rm results/*.dat
|
||||
|
||||
echo "Compiling binaries..."
|
||||
rm -f two_particles many_particles
|
||||
g++ -O3 -larmadillo -I ./include src/* two_particles.cpp -o two_particles
|
||||
g++ -O3 -larmadillo -I ./include src/* many_particles.cpp -o many_particles
|
||||
echo "Finished compiling."
|
||||
|
||||
if $PUEUE; then
|
||||
echo "Using pueue to generate results in parallel..."
|
||||
systemctl --user start pueued
|
||||
pueue parallel 8
|
||||
pueue clean
|
||||
COMMAND_PREFIX="pueue add "
|
||||
else
|
||||
echo "Generating results sequentially..."
|
||||
COMMAND_PREFIX=""
|
||||
fi
|
||||
|
||||
for N in 4000 8000 16000 32000
|
||||
do
|
||||
for interactions in "" "--disable-interactions"
|
||||
do
|
||||
for n in 1 2
|
||||
do
|
||||
$COMMAND_PREFIX ./two_particles -N $N -n $n $interactions
|
||||
$COMMAND_PREFIX ./two_particles -N $N -n $n $interactions -E
|
||||
$COMMAND_PREFIX ./two_particles -N $N -n $n $interactions -V
|
||||
done
|
||||
done
|
||||
# The analytical solution is only valid for one particle (second particle has v_z) without interactions
|
||||
$COMMAND_PREFIX ./two_particles -N $N -n 1 --disable-interactions -A
|
||||
done
|
||||
|
||||
if $NORMAL_MANY_PARTICLE_SIMULATION; then
|
||||
for N in 4000 8000 16000 32000
|
||||
do
|
||||
for interactions in "" "--disable-interactions"
|
||||
do
|
||||
for n in 10 50 100
|
||||
do
|
||||
$COMMAND_PREFIX ./many_particles -N $N -n $n $interactions
|
||||
$COMMAND_PREFIX ./many_particles -N $N -n $n $interactions -E
|
||||
$COMMAND_PREFIX ./many_particles -N $N -n $n $interactions -V
|
||||
done
|
||||
done
|
||||
done
|
||||
fi
|
||||
|
||||
if $OSCILLATING_FIELD_SIMULATION; then
|
||||
N=40000
|
||||
n=100
|
||||
reduced_output="--reduced-output"
|
||||
interactions="--disable-interactions"
|
||||
for amplitude in 0.1 0.4 0.7
|
||||
do
|
||||
for frequency in $(seq 0.2 0.02 2.5)
|
||||
do
|
||||
$COMMAND_PREFIX ./many_particles -N $N -t -n $n $interactions -O -f $amplitude -w $frequency $reduced_output
|
||||
$COMMAND_PREFIX ./many_particles -N $N -t -n $n $interactions -V -O -f $amplitude -w $frequency $reduced_output
|
||||
done
|
||||
done
|
||||
interactions=""
|
||||
for amplitude in 0.1 0.4 0.7
|
||||
do
|
||||
for frequency in $(seq 1.0 0.005 1.8)
|
||||
do
|
||||
$COMMAND_PREFIX ./many_particles -N $N -t -n $n $interactions -O -f $amplitude -w $frequency $reduced_output
|
||||
$COMMAND_PREFIX ./many_particles -N $N -t -n $n $interactions -V -O -f $amplitude -w $frequency $reduced_output
|
||||
done
|
||||
done
|
||||
fi
|
||||
|
||||
|
||||
echo "All tasks added to pueue."
|
||||
|
||||
echo "Use 'pueue status' to check the status of the tasks."
|
||||
|
||||
pueue status
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,85 @@
|
||||
#include <armadillo>
|
||||
#include <ostream>
|
||||
|
||||
#ifndef CLASSES_HPP
|
||||
#define CLASSES_HPP
|
||||
|
||||
class PenningTrap;
|
||||
|
||||
arma::vec standard_external_E_field(const arma::vec& r, double t, const PenningTrap& trap);
|
||||
arma::vec oscillating_external_E_field(const arma::vec& r, double t, const PenningTrap& trap);
|
||||
arma::vec standard_external_B_field(const arma::vec& r, double t, const PenningTrap& trap);
|
||||
arma::vec limited_range_external_B_field(const arma::vec& r, double t, const PenningTrap& trap);
|
||||
|
||||
class Particle {
|
||||
public:
|
||||
arma::vec position;
|
||||
arma::vec velocity;
|
||||
double mass;
|
||||
double charge;
|
||||
Particle(arma::vec pos, arma::vec vel, double m, double q);
|
||||
Particle(double m, double q);
|
||||
std::string info() const;
|
||||
bool inside_trap(double d) const;
|
||||
|
||||
};
|
||||
|
||||
class PenningTrap {
|
||||
private:
|
||||
std::vector<Particle> particles;
|
||||
bool coulomb_interaction = true;
|
||||
bool modified_fields = false;
|
||||
arma::vec (*_external_E_field)(const arma::vec& r, double t, const PenningTrap& trap) = standard_external_E_field;
|
||||
arma::vec (*_external_B_field)(const arma::vec& r, double t, const PenningTrap& trap) = standard_external_B_field;
|
||||
public:
|
||||
double B0;
|
||||
double V0;
|
||||
double d;
|
||||
double f = 0.0; // Amplitude of oscillating potential
|
||||
double omega_V = 0.0; // Angular frequency of oscillating potential
|
||||
PenningTrap(double B0_in, double V0_in, double d_in);
|
||||
void add_particle(Particle& p);
|
||||
void add_n_identical_particles(int n, double m, double q);
|
||||
arma::vec external_E_field(const arma::vec& r, double t) const {
|
||||
return _external_E_field(r, t, *this);
|
||||
}
|
||||
arma::vec external_B_field(const arma::vec& r, double t) const {
|
||||
return _external_B_field(r, t, *this);
|
||||
}
|
||||
void set_external_E_field(arma::vec (*E_field_func)(const arma::vec&, double, const PenningTrap&));
|
||||
void set_external_B_field(arma::vec (*B_field_func)(const arma::vec&, double, const PenningTrap&));
|
||||
arma::vec force_particle(int i, int j) const;
|
||||
arma::vec total_force_external(int i, double t) const;
|
||||
arma::vec total_force_particles(int i) const;
|
||||
arma::vec total_force(int i, double t) const;
|
||||
double potential_energy(int i) const;
|
||||
double total_potential_energy() const;
|
||||
double kinetic_energy(int i) const;
|
||||
double total_kinetic_energy() const;
|
||||
double total_energy() const;
|
||||
Particle& get_particle(int i);
|
||||
const Particle& get_particle(int i) const;
|
||||
int num_particles() const;
|
||||
void enable_interactions() {
|
||||
coulomb_interaction = true;
|
||||
}
|
||||
void disable_interactions() {
|
||||
coulomb_interaction = false;
|
||||
}
|
||||
bool get_coulomb_interaction() const {
|
||||
return coulomb_interaction;
|
||||
}
|
||||
bool get_modified_fields() const {
|
||||
return modified_fields;
|
||||
}
|
||||
std::string info() const;
|
||||
int num_inside_trap() const;
|
||||
};
|
||||
|
||||
|
||||
std::ostream& operator<<(std::ostream& os, const Particle& p);
|
||||
|
||||
std::ostream& operator<<(std::ostream& os, const PenningTrap& trap);
|
||||
|
||||
|
||||
#endif
|
||||
@@ -0,0 +1,13 @@
|
||||
#ifndef CONSTANTS_HPP
|
||||
#define CONSTANTS_HPP
|
||||
|
||||
namespace constants {
|
||||
const double k_e = 1.38935333e5; // Unit: u * um^3 / (us^2 * e^2), Coulomb's constant
|
||||
const double tesla = 9.64852558e1; // Conversion factor from tesla to u/(us * e)
|
||||
const double volt = 9.64852558e7; // Conversion factor from volt to u * um^2 / (us^2 * e)
|
||||
const double meter = 1e6; // Conversion factor from meter to micrometer
|
||||
const double second = 1e6; // Conversion factor from second to microsecond
|
||||
const double amu = 1.0; // Atomic mass unit in terms of itself (u)
|
||||
const double elementary_charge = 1.0; // Elementary charge in terms of itself (e)
|
||||
}
|
||||
#endif
|
||||
@@ -0,0 +1,92 @@
|
||||
#include "classes.hpp"
|
||||
#include <vector>
|
||||
#include <armadillo>
|
||||
|
||||
#ifndef SOLVERS_HPP
|
||||
#define SOLVERS_HPP
|
||||
|
||||
|
||||
arma::cube to_cube(const std::vector<arma::mat>& vec);
|
||||
|
||||
void save_to_file(const std::string& filename, const arma::cube& data);
|
||||
void save_to_file(const std::string& filename, const arma::vec& data);
|
||||
|
||||
|
||||
class Solver {
|
||||
protected:
|
||||
PenningTrap& trap;
|
||||
double dt;
|
||||
int steps;
|
||||
std::vector<arma::mat> positions;
|
||||
std::vector<arma::mat> velocities;
|
||||
std::vector<int> inside_trap_count;
|
||||
std::vector<double> potential_energies;
|
||||
std::vector<double> kinetic_energies;
|
||||
std::vector<double> total_energies;
|
||||
virtual void step() = 0;
|
||||
void record_position();
|
||||
void record_velocity();
|
||||
void record_energies();
|
||||
void record_trap_count();
|
||||
public:
|
||||
Solver(PenningTrap& trap_in, double dt_in);
|
||||
void simulate(int num_steps);
|
||||
void simulate_time(double total_time) {
|
||||
int num_steps = total_steps(total_time);
|
||||
simulate(num_steps);
|
||||
}
|
||||
std::vector<arma::mat> get_positions();
|
||||
std::vector<arma::mat> get_velocities();
|
||||
std::vector<double> get_potential_energies();
|
||||
std::vector<double> get_kinetic_energies();
|
||||
std::vector<double> get_total_energies();
|
||||
std::vector<int> get_inside_trap_count();
|
||||
int get_current_step();
|
||||
int total_steps(double total_time);
|
||||
void save(std::string base_filename);
|
||||
void save_counts(std::string base_filename);
|
||||
void record_current_step() {
|
||||
record_position();
|
||||
record_velocity();
|
||||
if (!trap.get_modified_fields()) {
|
||||
record_energies();
|
||||
}
|
||||
record_trap_count();
|
||||
}
|
||||
bool record_positions_and_velocities = true;
|
||||
bool record_inside_trap_count = true;
|
||||
};
|
||||
|
||||
class EulerSolver : public Solver {
|
||||
public:
|
||||
EulerSolver(PenningTrap& trap_in, double dt_in);
|
||||
void step() override;
|
||||
};
|
||||
class RK4Solver : public Solver {
|
||||
public:
|
||||
RK4Solver(PenningTrap& trap_in, double dt_in);
|
||||
void step() override;
|
||||
};
|
||||
class VelocityVerletSolver : public Solver {
|
||||
private:
|
||||
std::vector<arma::vec> accelerations;
|
||||
public:
|
||||
VelocityVerletSolver(PenningTrap& trap_in, double dt_in);
|
||||
void step() override;
|
||||
};
|
||||
|
||||
class AnalyticalSolver : public Solver {
|
||||
private:
|
||||
std::vector<double> w_0;
|
||||
std::vector<double> w_z;
|
||||
std::vector<double> w_plus;
|
||||
std::vector<double> w_minus;
|
||||
std::vector<double> A_plus;
|
||||
std::vector<double> A_minus;
|
||||
std::vector<double> z_0;
|
||||
public:
|
||||
AnalyticalSolver(PenningTrap& trap_in, double dt_in);
|
||||
void step() override;
|
||||
};
|
||||
|
||||
#endif
|
||||
@@ -0,0 +1,173 @@
|
||||
#include <armadillo>
|
||||
#include "classes.hpp"
|
||||
#include "solvers.hpp"
|
||||
#include "constants.hpp"
|
||||
#include "argparse/argparse.hpp"
|
||||
|
||||
#include <iostream>
|
||||
#include <ctime>
|
||||
|
||||
using namespace std;
|
||||
|
||||
void print_first_particles(PenningTrap& trap, int num_particles) {
|
||||
for (int i = 0; i < num_particles && i < trap.num_particles(); ++i) {
|
||||
cout << "Particle " << i + 1 << ": " << endl << trap.get_particle(i) << endl;
|
||||
}
|
||||
}
|
||||
|
||||
void save_timing(clock_t start, clock_t end, const string& filename) {
|
||||
double cpu_time_used = ((double)(end - start)) / CLOCKS_PER_SEC;
|
||||
ofstream file(filename);
|
||||
if (file.is_open()) {
|
||||
file << "CPU time used: " << cpu_time_used << " seconds" << endl;
|
||||
file.close();
|
||||
} else {
|
||||
cerr << "Unable to open file for writing timing information." << endl;
|
||||
}
|
||||
}
|
||||
|
||||
int main(int argc, char* argv[]) {
|
||||
argparse::ArgumentParser program("many_particles");
|
||||
program.add_argument("-v", "--verbose")
|
||||
.help("Enable verbose output of first particles")
|
||||
.flag();
|
||||
program.add_argument("-N", "--steps")
|
||||
.help("Number of time steps")
|
||||
.default_value(4000)
|
||||
.scan<'i', int>();
|
||||
program.add_argument("-n", "--num_particles")
|
||||
.help("Number of particles")
|
||||
.default_value(100)
|
||||
.scan<'i', int>();
|
||||
program.add_argument("-i", "--disable-interactions")
|
||||
.help("Disable Coulomb interactions")
|
||||
.default_value(true)
|
||||
.implicit_value(false);
|
||||
program.add_argument("-O", "--oscillating-potential")
|
||||
.help("Enable oscillating potential")
|
||||
.flag();
|
||||
program.add_argument("-f", "--amplitude")
|
||||
.help("Amplitude of oscillating potential")
|
||||
.default_value(0.0)
|
||||
.scan<'f', double>();
|
||||
program.add_argument("-w", "--frequency")
|
||||
.help("Angular frequency of oscillating potential")
|
||||
.default_value(0.0)
|
||||
.scan<'f', double>();
|
||||
program.add_argument("-r", "--reduced-output")
|
||||
.help("Reduce output for oscillating potential to only final count")
|
||||
.flag();
|
||||
program.add_argument("-E", "--euler")
|
||||
.help("Use Euler solver instead of RK4")
|
||||
.flag();
|
||||
program.add_argument("-V", "--velocity-verlet")
|
||||
.help("Use Velocity Verlet solver instead of RK4")
|
||||
.flag();
|
||||
program.add_argument("-t", "--timing")
|
||||
.help("Enable timing of the simulation")
|
||||
.flag();
|
||||
|
||||
try {
|
||||
program.parse_args(argc, argv);
|
||||
} catch (const std::runtime_error& err) {
|
||||
std::cerr << err.what() << std::endl;
|
||||
std::cerr << program.help().str() << std::endl;
|
||||
exit(1);
|
||||
}
|
||||
|
||||
bool verbose = program.get<bool>("--verbose");
|
||||
|
||||
int N = program.get<int>("--steps");
|
||||
int n_particles = program.get<int>("--num_particles");
|
||||
bool interactions = program.get<bool>("--disable-interactions");
|
||||
bool oscillating_potential = program.get<bool>("--oscillating-potential");
|
||||
bool use_euler = program.get<bool>("--euler");
|
||||
bool use_velocity_verlet = program.get<bool>("--velocity-verlet");
|
||||
bool timing = program.get<bool>("--timing");
|
||||
|
||||
double B_0 = 1.0; // Tesla
|
||||
double V_0 = 0.025; // Volt
|
||||
double d = 500e-6; // meter
|
||||
|
||||
PenningTrap trap(B_0, V_0, d);
|
||||
if (interactions) {
|
||||
trap.enable_interactions();
|
||||
}
|
||||
else {
|
||||
trap.disable_interactions();
|
||||
}
|
||||
|
||||
|
||||
bool reduced_output = program.get<bool>("--reduced-output"); // Only relevant if oscillating potential is enabled
|
||||
if (oscillating_potential) {
|
||||
trap.set_external_E_field(oscillating_external_E_field);
|
||||
trap.set_external_B_field(limited_range_external_B_field);
|
||||
trap.f = program.get<double>("--amplitude");
|
||||
trap.omega_V = program.get<double>("--frequency");
|
||||
}
|
||||
|
||||
trap.add_n_identical_particles(n_particles, 40.078 * constants::amu, 1 * constants::elementary_charge);
|
||||
if (verbose) {
|
||||
cout << trap.info() << endl;
|
||||
print_first_particles(trap, 10);
|
||||
}
|
||||
|
||||
|
||||
double dt = 50e-6 / N; // seconds
|
||||
|
||||
|
||||
unique_ptr<Solver> solver;
|
||||
string solver_name;
|
||||
|
||||
if (use_euler) {
|
||||
solver = make_unique<EulerSolver>(trap, dt);
|
||||
solver_name = "Euler";
|
||||
} else if (use_velocity_verlet) {
|
||||
solver = make_unique<VelocityVerletSolver>(trap, dt);
|
||||
solver_name = "VelocityVerlet";
|
||||
} else {
|
||||
solver = make_unique<RK4Solver>(trap, dt);
|
||||
solver_name = "RK4";
|
||||
}
|
||||
clock_t start, end;
|
||||
start = clock();
|
||||
if (oscillating_potential) {
|
||||
solver->record_positions_and_velocities = false; // Disable recording for performance
|
||||
if (reduced_output) {
|
||||
solver->record_inside_trap_count = false; // Disable recording for performance
|
||||
}
|
||||
solver->simulate(N);
|
||||
end = clock();
|
||||
string folder = "results/" + solver_name + "/many_particles_" + to_string(trap.num_particles()) + "_particles_" + to_string(int(N)) + "_steps_" + (interactions ? "enabled" : "disabled") + "_interactions_oscillating_potential";
|
||||
string filename = "f" + to_string(trap.f) + "_w" + to_string(trap.omega_V);
|
||||
filesystem::create_directories(folder);
|
||||
|
||||
if (reduced_output) {
|
||||
solver->record_current_step(); // Record only final step
|
||||
}
|
||||
solver->save_counts(folder + "/" + filename);
|
||||
cout << "Final number of particles inside trap: " << trap.num_inside_trap() << " out of " << trap.num_particles() << endl;
|
||||
if (timing) {
|
||||
save_timing(start, end, folder + "/" + filename + "_timing.txt");
|
||||
}
|
||||
return 0;
|
||||
|
||||
}
|
||||
solver->simulate(N);
|
||||
end = clock();
|
||||
string folder = "results/" + solver_name;
|
||||
filesystem::create_directories(folder);
|
||||
string filename = "/many_particles_" + to_string(trap.num_particles()) + "_particles_" + to_string(int(N)) + "_steps_" + (interactions ? "enabled" : "disabled") + "_interactions";
|
||||
solver->save(folder + filename);
|
||||
if (verbose) {
|
||||
vector<arma::mat> positions = solver->get_positions();
|
||||
cout << "Final positions of particles: " << endl;
|
||||
print_first_particles(trap, 10);
|
||||
}
|
||||
if (timing) {
|
||||
save_timing(start, end, folder + filename + "_timing.txt");
|
||||
}
|
||||
|
||||
return 0;
|
||||
}
|
||||
|
||||
Executable
+36
@@ -0,0 +1,36 @@
|
||||
#!/usr/bin/env bash
|
||||
|
||||
interval=5 # seconds between updates
|
||||
|
||||
while true; do
|
||||
data=$(pueue status --json)
|
||||
|
||||
queued=$(jq '.tasks | map(select(.status=="Queued")) | length' <<< "$data")
|
||||
running=$(jq '.tasks | map(select(.status=="Running")) | length' <<< "$data")
|
||||
total=$(jq '.tasks | length' <<< "$data")
|
||||
|
||||
# average duration of finished tasks (seconds)
|
||||
avg=$(jq '.tasks
|
||||
| map(select(.start != null and .end != null)
|
||||
| ((.end[0:19] + "Z") | fromdate) - ((.start[0:19] + "Z") | fromdate))
|
||||
| if length > 0 then add / length else 0 end' <<< "$data")
|
||||
|
||||
clear
|
||||
echo "Queued: $queued"
|
||||
echo "Running: $running"
|
||||
echo "Total: $total"
|
||||
|
||||
# Only calculate ETA if avg > 0
|
||||
if (( $(echo "$avg > 0" | bc -l) )); then
|
||||
remaining=$((queued + running))
|
||||
eta_seconds=$(echo "$remaining * $avg" | bc -l)
|
||||
finish_time=$(date -d "@$(($(date +%s) + ${eta_seconds%.*}))" +"%Y-%m-%d %H:%M:%S")
|
||||
printf "Average duration: %.2fs\n" "$avg"
|
||||
printf "ETA: ~%.0fs (finishing around %s)\n" "$eta_seconds" "$finish_time"
|
||||
else
|
||||
echo "ETA: Not enough finished jobs in history yet."
|
||||
fi
|
||||
|
||||
sleep $interval
|
||||
done
|
||||
|
||||
@@ -0,0 +1,35 @@
|
||||
import matplotlib.pyplot as plt
|
||||
import numpy as np
|
||||
import os
|
||||
import typer
|
||||
|
||||
app = typer.Typer()
|
||||
|
||||
def load_data(total_energy_file):
|
||||
total_energy = np.loadtxt(total_energy_file)
|
||||
kin_file = total_energy_file.replace("total_energies", "kinetic_energies")
|
||||
pot_file = total_energy_file.replace("total_energies", "potential_energies")
|
||||
kinetic_energy = np.loadtxt(kin_file)
|
||||
potential_energy = np.loadtxt(pot_file)
|
||||
step_count = len(total_energy)
|
||||
return total_energy, kinetic_energy, potential_energy, step_count
|
||||
|
||||
@app.command()
|
||||
def plot_kinetic_energy(energy_files: list[str]=[], labels: list[str]=[], mass: float=1.0):
|
||||
"""Plot kinetic energy from velocity files."""
|
||||
for i, energy_file in enumerate(energy_files):
|
||||
total_energy, kinetic_energy, potential_energy, step_count = load_data(energy_file)
|
||||
time = np.linspace(0, 50e-6, step_count)
|
||||
color = f"C{i % 10}"
|
||||
plt.plot(time, total_energy, label=labels[i], color=color, linestyle="-")
|
||||
plt.plot(time, potential_energy, color=color, linestyle="--")
|
||||
plt.plot(time, kinetic_energy, color=color, linestyle=":")
|
||||
plt.xlabel('Time (s)')
|
||||
plt.ylabel('Energy (J)')
|
||||
plt.title('Energy Evolution')
|
||||
plt.legend()
|
||||
plt.show()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
app()
|
||||
@@ -0,0 +1,49 @@
|
||||
import matplotlib.pyplot as plt
|
||||
import numpy as np
|
||||
import os
|
||||
import argparse
|
||||
|
||||
argparser = argparse.ArgumentParser()
|
||||
argparser.add_argument("directory", type=str, help="Directory containing the data files")
|
||||
args = argparser.parse_args()
|
||||
directory = args.directory
|
||||
|
||||
def get_filenames(directory):
|
||||
filenames = [os.path.join(directory, f) for f in os.listdir(directory) if f.endswith(".red")]
|
||||
return filenames
|
||||
|
||||
def get_pars(filename):
|
||||
amplitude = float(filename.split("_")[0][1:])
|
||||
omega = float(filename.split("_")[1][1:])
|
||||
return amplitude, omega
|
||||
|
||||
def load_data(filename):
|
||||
with open(filename, "r") as file:
|
||||
x = int(file.readline().split()[1])
|
||||
return x
|
||||
|
||||
def plot_resonances(amplitudes, omegas, counts):
|
||||
plt.figure(figsize=(8, 6))
|
||||
unique_amps = sorted(set(amplitudes))
|
||||
for amp in unique_amps:
|
||||
mask = np.array(amplitudes) == amp
|
||||
omegas_subset = np.array(omegas)[mask]
|
||||
counts_subset = np.array(counts)[mask]
|
||||
omegas_sorted, counts_sorted = zip(*sorted(zip(omegas_subset, counts_subset)))
|
||||
plt.plot(omegas_sorted, counts_sorted, label=f'Amplitude {amp}')
|
||||
plt.xlabel('Omega')
|
||||
plt.ylabel('Count inside trap')
|
||||
plt.legend()
|
||||
plt.show()
|
||||
|
||||
if __name__ == "__main__":
|
||||
filenames = get_filenames(directory)
|
||||
amplitudes, omegas, counts = [], [], []
|
||||
for filename in filenames:
|
||||
amplitude, omega = get_pars(os.path.basename(filename))
|
||||
count = load_data(filename)
|
||||
amplitudes.append(amplitude)
|
||||
omegas.append(omega)
|
||||
counts.append(count)
|
||||
plot_resonances(amplitudes, omegas, counts)
|
||||
|
||||
@@ -0,0 +1,93 @@
|
||||
import matplotlib.pyplot as plt
|
||||
import numpy as np
|
||||
import os
|
||||
import typer
|
||||
from typing import Literal
|
||||
|
||||
app = typer.Typer()
|
||||
|
||||
def load_trajectory_data(filename):
|
||||
data = np.loadtxt(filename)
|
||||
step_count = int(float(filename.split("/")[-1].split("_")[5]))
|
||||
particle_count = int(float(filename.split("/")[-1].split("_")[2]))
|
||||
single_particle = particle_count == 1
|
||||
if single_particle:
|
||||
data = data.reshape((data.shape[0], 1))
|
||||
data = data.reshape(step_count, 3, particle_count).transpose(1, 2, 0)
|
||||
return data, step_count, particle_count
|
||||
|
||||
@app.command()
|
||||
def plot_z_evolution(trajectory_file: str):
|
||||
"""Plot the z-coordinate evolution from a trajectory file."""
|
||||
data, step_count, particle_count = load_trajectory_data(trajectory_file)
|
||||
time = np.linspace(0, 50e-6, step_count)
|
||||
for i in range(particle_count):
|
||||
plt.plot(time, data[2, i, :] * 1e-6, label=f'Particle {i+1}')
|
||||
plt.xlabel('Time (s)')
|
||||
plt.ylabel('Z Position (m)')
|
||||
plt.title('Z Coordinate Evolution')
|
||||
plt.legend()
|
||||
plt.show()
|
||||
|
||||
@app.command()
|
||||
def plot_xy_trajectory(trajectory_file: str, plot_start_end: bool = True, plot_time: bool = False):
|
||||
"""Plot the XY trajectory from a trajectory file."""
|
||||
data, step_count, particle_count = load_trajectory_data(trajectory_file)
|
||||
for i in range(particle_count):
|
||||
if plot_start_end:
|
||||
plt.scatter(data[0, i, 0] * 1e-6, data[1, i, 0] * 1e-6, marker="o", color="black", label="Start" if i == 0 else "")
|
||||
plt.scatter(data[0, i, -1] * 1e-6, data[1, i, -1] * 1e-6, marker="x", color="red", label="End" if i == 0 else "")
|
||||
if plot_time:
|
||||
marker = [".", ",", "o", "v", "^", "<", ">", "1", "2", "3", "4", "8"][i % 12]
|
||||
cm = plt.get_cmap('viridis')
|
||||
c = cm(np.linspace(0, 1, step_count))
|
||||
plt.scatter(data[0, i, :] * 1e-6, data[1, i, :] * 1e-6, label=f'Particle {i+1}', marker=marker, c=c)
|
||||
else:
|
||||
ls = "-"
|
||||
c = f"C{i % 10}"
|
||||
plt.plot(data[0, i, :] * 1e-6, data[1, i, :] * 1e-6, label=f'Particle {i+1}', linestyle=ls, color=c)
|
||||
plt.xlabel('X Position (m)')
|
||||
plt.ylabel('Y Position (m)')
|
||||
plt.title('XY Trajectory')
|
||||
plt.axis('equal')
|
||||
if plot_time:
|
||||
#sm = plt.cm.ScalarMappable(cmap=cm, norm=plt.Normalize(vmin=0, vmax=50e-6))
|
||||
#sm.set_array([])
|
||||
#cbar = plt.colorbar(sm)
|
||||
#cbar.set_label('Time')
|
||||
pass
|
||||
plt.legend()
|
||||
plt.show()
|
||||
|
||||
@app.command()
|
||||
def phase_space_plot(trajectory_file: str, velocity_file: str, axis: Literal['x', 'y', 'z'] = 'z'):
|
||||
"""Plot phase space (Z vs Vz) from trajectory and velocity files."""
|
||||
axis_index = {'x': 0, 'y': 1, 'z': 2}[axis]
|
||||
traj_data, step_count, particle_count = load_trajectory_data(trajectory_file)
|
||||
vel_data, _, _ = load_trajectory_data(velocity_file)
|
||||
for i in range(particle_count):
|
||||
plt.plot(traj_data[axis_index, i, :] * 1e-6, vel_data[axis_index, i, :] * 1e-6, label=f'Particle {i+1}')
|
||||
plt.xlabel(f'{axis.upper()} Position (m)')
|
||||
plt.ylabel(f'{axis.upper()} Velocity (m/s)')
|
||||
plt.title(f'Phase Space Plot ({axis.upper()} vs V{axis.upper()})')
|
||||
plt.legend()
|
||||
plt.show()
|
||||
|
||||
@app.command()
|
||||
def plot_error(trajectory_files: list[str], analytical_file: str):
|
||||
"""Plot error between numerical and analytical trajectories."""
|
||||
analytical_data, step_count, particle_count = load_trajectory_data(analytical_file)
|
||||
time = np.linspace(0, 50e-6, step_count)
|
||||
for traj_file in trajectory_files:
|
||||
numerical_data, _, _ = load_trajectory_data(traj_file)
|
||||
error = numerical_data - analytical_data
|
||||
relative_error = np.mean(np.linalg.norm(error, axis=0) / np.linalg.norm(analytical_data, axis=0), axis=0)
|
||||
plt.plot(time, relative_error, label=os.path.basename(traj_file))
|
||||
plt.xlabel('Time (s)')
|
||||
plt.ylabel('Relative Error')
|
||||
plt.title('Relative Error between Numerical and Analytical Trajectories')
|
||||
plt.legend()
|
||||
plt.show()
|
||||
|
||||
if __name__ == "__main__":
|
||||
app()
|
||||
@@ -0,0 +1,223 @@
|
||||
#include "classes.hpp"
|
||||
#include "constants.hpp"
|
||||
#include <iostream>
|
||||
|
||||
// Initialize particle with position and velocity
|
||||
// Mass in atomic mass units (u), charge in elementary charges (e)
|
||||
// Positions in meters, velocities in meters/second
|
||||
Particle::Particle(arma::vec pos, arma::vec vel, double m, double q)
|
||||
: position(pos * constants::meter), velocity(vel * constants::meter / constants::second), mass(m * constants::amu), charge(q * constants::elementary_charge) {}
|
||||
|
||||
// Initialize particle at origin with zero velocity
|
||||
Particle::Particle(double m, double q)
|
||||
: position(arma::vec({0,0,0})), velocity(arma::vec({0,0,0})), mass(m * constants::amu), charge(q * constants::elementary_charge) {}
|
||||
|
||||
std::string Particle::info() const {
|
||||
std::string info_str = "Particle (m = " + std::to_string(mass) + " u, q = " + std::to_string(charge) + " e)\n";
|
||||
info_str += "Position (um): [" + std::to_string(position(0)) + ", " + std::to_string(position(1)) + ", " + std::to_string(position(2)) + "]\n";
|
||||
info_str += "Velocity (um/us): [" + std::to_string(velocity(0)) + ", " + std::to_string(velocity(1)) + ", " + std::to_string(velocity(2)) + "]\n";
|
||||
return info_str;
|
||||
}
|
||||
|
||||
bool Particle::inside_trap(double d) const {
|
||||
return arma::norm(position) <= d;
|
||||
}
|
||||
|
||||
// Empty Penning Trap with specified field parameters
|
||||
// B0 in tesla, V0 in volt, d in meters
|
||||
PenningTrap::PenningTrap(double B0_in, double V0_in, double d_in)
|
||||
: B0(B0_in * constants::tesla), V0(V0_in * constants::volt), d(d_in * constants::meter) {}
|
||||
|
||||
// Add particle to trap
|
||||
void PenningTrap::add_particle(Particle& p) {
|
||||
particles.push_back(p);
|
||||
}
|
||||
|
||||
// Add n identical particles with mass m and charge q
|
||||
void PenningTrap::add_n_identical_particles(int n, double m, double q) {
|
||||
for (int i = 0; i < n; ++i) {
|
||||
arma::vec pos = arma::vec(3).randn() * 0.1 * d; // Random position within 10% of d (unit: um)
|
||||
arma::vec vel = arma::vec(3).randn() * 0.1 * d; // Random velocity within 10% of d/s (unit: um/us)
|
||||
Particle p(pos / constants::meter, vel, m, q);
|
||||
add_particle(p);
|
||||
}
|
||||
}
|
||||
|
||||
void PenningTrap::set_external_E_field(arma::vec (*E_field_func)(const arma::vec&, double, const PenningTrap&)) {
|
||||
_external_E_field = E_field_func;
|
||||
modified_fields = true;
|
||||
}
|
||||
|
||||
void PenningTrap::set_external_B_field(arma::vec (*B_field_func)(const arma::vec&, double, const PenningTrap&)) {
|
||||
_external_B_field = B_field_func;
|
||||
modified_fields = true;
|
||||
}
|
||||
|
||||
|
||||
// Force on particle i from particle j due to Coulomb interaction
|
||||
// Neglecting magnetic fields from moving charges
|
||||
arma::vec PenningTrap::force_particle(int i, int j) const {
|
||||
if (!coulomb_interaction) {
|
||||
return arma::vec({0, 0, 0});
|
||||
}
|
||||
arma::vec r_i = particles[i].position;
|
||||
arma::vec r_j = particles[j].position;
|
||||
arma::vec r_ij = r_i - r_j;
|
||||
double distance = arma::norm(r_ij);
|
||||
if (distance == 0) {
|
||||
return arma::vec({0, 0, 0}); // Avoid division by zero
|
||||
}
|
||||
double force_magnitude_per_dist = (constants::k_e * particles[i].charge * particles[j].charge) / (distance * distance * distance);
|
||||
return force_magnitude_per_dist * r_ij;
|
||||
}
|
||||
|
||||
// Total force on particle i from external fields
|
||||
// Sum of electric and magnetic forces
|
||||
arma::vec PenningTrap::total_force_external(int i, double t) const {
|
||||
arma::vec E = external_E_field(get_particle(i).position, t);
|
||||
arma::vec B = external_B_field(get_particle(i).position, t);
|
||||
arma::vec v = get_particle(i).velocity;
|
||||
arma::vec F_electric = get_particle(i).charge * E;
|
||||
arma::vec F_magnetic = get_particle(i).charge * arma::cross(v, B);
|
||||
return F_electric + F_magnetic;
|
||||
}
|
||||
|
||||
// Total force on particle i from all other particles
|
||||
arma::vec PenningTrap::total_force_particles(int i) const {
|
||||
arma::vec total_force = arma::vec({0, 0, 0});
|
||||
if (!coulomb_interaction) { // Add the check here as well to avoid unnecessary calls to force_particle
|
||||
return total_force;
|
||||
}
|
||||
for (size_t j = 0; j < particles.size(); ++j) {
|
||||
if (j != i) {
|
||||
total_force += force_particle(i, j);
|
||||
}
|
||||
}
|
||||
return total_force;
|
||||
}
|
||||
|
||||
// Total force on particle i from all other particles and external fields
|
||||
arma::vec PenningTrap::total_force(int i, double t) const {
|
||||
return total_force_external(i, t) + total_force_particles(i);
|
||||
}
|
||||
|
||||
// Potential energy of particle i in the trap
|
||||
double PenningTrap::potential_energy(int i) const {
|
||||
if (modified_fields) {
|
||||
throw std::runtime_error("Potential energy calculation not valid with modified external fields.");
|
||||
}
|
||||
double q = particles[i].charge;
|
||||
arma::vec r = particles[i].position;
|
||||
double V = (V0 / (d * d)) * (r(2) * r(2) - 0.5 * (r(0) * r(0) + r(1) * r(1)));
|
||||
if (coulomb_interaction) {
|
||||
// Add contributions from other particles
|
||||
for (size_t j = 0; j < particles.size(); ++j) {
|
||||
if (j != i) {
|
||||
arma::vec r_j = particles[j].position;
|
||||
double distance = arma::norm(r - r_j);
|
||||
if (distance != 0) {
|
||||
V += (constants::k_e * particles[j].charge) / distance;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
return q * V;
|
||||
}
|
||||
|
||||
// Total potential energy of all particles in the trap
|
||||
double PenningTrap::total_potential_energy() const {
|
||||
if (modified_fields) {
|
||||
throw std::runtime_error("Potential energy calculation not valid with modified external fields.");
|
||||
}
|
||||
double total_U = 0.0;
|
||||
for (size_t i = 0; i < particles.size(); ++i) {
|
||||
total_U += potential_energy(i);
|
||||
}
|
||||
return total_U;
|
||||
}
|
||||
|
||||
// Kinetic energy of particle i
|
||||
double PenningTrap::kinetic_energy(int i) const {
|
||||
double m = particles[i].mass;
|
||||
arma::vec v = particles[i].velocity;
|
||||
return 0.5 * m * arma::dot(v, v);
|
||||
}
|
||||
// Total kinetic energy of all particles in the trap
|
||||
double PenningTrap::total_kinetic_energy() const {
|
||||
double total_K = 0.0;
|
||||
for (size_t i = 0; i < particles.size(); ++i) {
|
||||
total_K += kinetic_energy(i);
|
||||
}
|
||||
return total_K;
|
||||
}
|
||||
|
||||
// Total energy of all particles in the trap
|
||||
double PenningTrap::total_energy() const {
|
||||
return total_kinetic_energy() + total_potential_energy();
|
||||
}
|
||||
|
||||
|
||||
// Getter function for particles
|
||||
Particle& PenningTrap::get_particle(int i) {
|
||||
return particles[i];
|
||||
}
|
||||
|
||||
const Particle& PenningTrap::get_particle(int i) const {
|
||||
return particles[i];
|
||||
}
|
||||
|
||||
// Number of particles in trap
|
||||
int PenningTrap::num_particles() const{
|
||||
return particles.size();
|
||||
}
|
||||
|
||||
|
||||
std::string PenningTrap::info() const {
|
||||
std::string info_str = "Penning Trap (B0 = " + std::to_string(B0) + " u/(us*e), V0 = " + std::to_string(V0) + " u*um^2/(us^2*e), d = " + std::to_string(d) + " um)\n";
|
||||
info_str += "Number of particles: " + std::to_string(num_particles()) + "\n";
|
||||
info_str += "Coulomb interactions: " + std::string(coulomb_interaction ? "enabled" : "disabled") + "\n";
|
||||
return info_str;
|
||||
}
|
||||
|
||||
int PenningTrap::num_inside_trap() const {
|
||||
int count = 0;
|
||||
for (int i = 0; i < num_particles(); ++i) {
|
||||
if (particles[i].inside_trap(d)) {
|
||||
count++;
|
||||
}
|
||||
}
|
||||
return count;
|
||||
}
|
||||
|
||||
// Electric field at position r and time t
|
||||
// Electric field from V = V0 * (z^2 - (x^2 + y^2)/2) / d^2
|
||||
arma::vec standard_external_E_field(const arma::vec& r, double t, const PenningTrap& trap) {
|
||||
return (trap.V0 / (trap.d * trap.d)) * arma::vec({r(0), r(1), -2 * r(2)});
|
||||
}
|
||||
|
||||
// Magnetic field at position r and time t
|
||||
// Constant magnetic field in z-direction
|
||||
arma::vec standard_external_B_field(const arma::vec& r, double t, const PenningTrap& trap) {
|
||||
return arma::vec({0, 0, trap.B0});
|
||||
}
|
||||
|
||||
arma::vec oscillating_external_E_field(const arma::vec& r, double t, const PenningTrap& trap) {
|
||||
if (arma::norm(r) > trap.d) {
|
||||
return arma::vec({0, 0, 0}); // No electric field outside the trap
|
||||
}
|
||||
return standard_external_E_field(r, t, trap) * (1 + trap.f * cos(trap.omega_V * t));
|
||||
}
|
||||
|
||||
arma::vec limited_range_external_B_field(const arma::vec& r, double t, const PenningTrap& trap) {
|
||||
if (arma::norm(r) > trap.d) {
|
||||
return arma::vec({0, 0, 0}); // No magnetic field outside the trap
|
||||
}
|
||||
return standard_external_B_field(r, t, trap);
|
||||
}
|
||||
|
||||
std::ostream& operator<<(std::ostream& os, const Particle& p) {
|
||||
return os << p.info();
|
||||
}
|
||||
std::ostream& operator<<(std::ostream& os, const PenningTrap& trap) {
|
||||
return os << trap.info();
|
||||
}
|
||||
@@ -0,0 +1,340 @@
|
||||
#include "classes.hpp"
|
||||
#include "solvers.hpp"
|
||||
#include "constants.hpp"
|
||||
|
||||
|
||||
// Initialize base solver class
|
||||
// dt in seconds
|
||||
Solver::Solver(PenningTrap& trap_in, double dt_in)
|
||||
: trap(trap_in), dt(dt_in * constants::second), steps(0) {}
|
||||
|
||||
// Simulate the system for a given number of steps
|
||||
void Solver::simulate(int num_steps) {
|
||||
for (int i = 0; i < num_steps; i++) {
|
||||
step();
|
||||
if (record_positions_and_velocities) { // Allow disabling of recording for performance
|
||||
// Store positions of all particles after each step
|
||||
record_position();
|
||||
// Store velocities of all particles after each step
|
||||
record_velocity();
|
||||
if (!trap.get_modified_fields()) {
|
||||
record_energies();
|
||||
}
|
||||
}
|
||||
if (record_inside_trap_count) {
|
||||
record_trap_count();
|
||||
}
|
||||
steps++;
|
||||
}
|
||||
}
|
||||
|
||||
void Solver::record_position() {
|
||||
arma::mat current_positions(3, trap.num_particles());
|
||||
for (int j = 0; j < trap.num_particles(); j++) {
|
||||
current_positions.col(j) = trap.get_particle(j).position;
|
||||
}
|
||||
positions.push_back(current_positions);
|
||||
}
|
||||
void Solver::record_velocity() {
|
||||
arma::mat current_velocities(3, trap.num_particles());
|
||||
for (int j = 0; j < trap.num_particles(); j++) {
|
||||
current_velocities.col(j) = trap.get_particle(j).velocity;
|
||||
}
|
||||
velocities.push_back(current_velocities);
|
||||
}
|
||||
|
||||
void Solver::record_energies() {
|
||||
double U = trap.total_potential_energy();
|
||||
double K = trap.total_kinetic_energy();
|
||||
double E = trap.total_energy();
|
||||
potential_energies.push_back(U);
|
||||
kinetic_energies.push_back(K);
|
||||
total_energies.push_back(E);
|
||||
}
|
||||
|
||||
void Solver::record_trap_count() {
|
||||
inside_trap_count.push_back(trap.num_inside_trap());
|
||||
}
|
||||
|
||||
// Get recorded positions of particles
|
||||
std::vector<arma::mat> Solver::get_positions() {
|
||||
return positions;
|
||||
}
|
||||
|
||||
// Get recorded velocities of particles
|
||||
std::vector<arma::mat> Solver::get_velocities() {
|
||||
return velocities;
|
||||
}
|
||||
std::vector<double> Solver::get_potential_energies() {
|
||||
return potential_energies;
|
||||
}
|
||||
std::vector<double> Solver::get_kinetic_energies() {
|
||||
return kinetic_energies;
|
||||
}
|
||||
std::vector<double> Solver::get_total_energies() {
|
||||
return total_energies;
|
||||
}
|
||||
|
||||
std::vector<int> Solver::get_inside_trap_count() {
|
||||
return inside_trap_count;
|
||||
}
|
||||
|
||||
// Get the current step count
|
||||
int Solver::get_current_step() {
|
||||
return steps;
|
||||
}
|
||||
|
||||
// Calculate total steps for a given total time
|
||||
// total_time in seconds
|
||||
int Solver::total_steps(double total_time) {
|
||||
return static_cast<int>(total_time * constants::second / dt);
|
||||
}
|
||||
|
||||
void Solver::save(std::string base_filename) {
|
||||
arma::cube pos_cube = to_cube(positions);
|
||||
save_to_file(base_filename + "_positions.dat", pos_cube);
|
||||
arma::cube vel_cube = to_cube(velocities);
|
||||
save_to_file(base_filename + "_velocities.dat", vel_cube);
|
||||
if (!trap.get_modified_fields()) {
|
||||
arma::vec U = arma::vec(potential_energies);
|
||||
save_to_file(base_filename + "_potential_energies.dat", U);
|
||||
arma::vec K = arma::vec(kinetic_energies);
|
||||
save_to_file(base_filename + "_kinetic_energies.dat", K);
|
||||
arma::vec E = arma::vec(total_energies);
|
||||
save_to_file(base_filename + "_total_energies.dat", E);
|
||||
}
|
||||
}
|
||||
|
||||
void Solver::save_counts(std::string base_filename) {
|
||||
std::string extension = ".dat";
|
||||
if (inside_trap_count.size() == 1) {
|
||||
// Reduced to last element if only one entry
|
||||
extension = ".red";
|
||||
}
|
||||
std::ofstream file(base_filename + "_inside_trap_count" + extension);
|
||||
for (size_t i = 0; i < inside_trap_count.size(); i++) {
|
||||
file << i << " " << inside_trap_count[i] << "\n";
|
||||
}
|
||||
file.close();
|
||||
}
|
||||
|
||||
// Initialize Euler solver
|
||||
EulerSolver::EulerSolver(PenningTrap& trap_in, double dt_in)
|
||||
: Solver(trap_in, dt_in) {}
|
||||
|
||||
// Perform a single Euler integration step
|
||||
void EulerSolver::step() {
|
||||
int n = trap.num_particles();
|
||||
for (int i = 0; i < n; i++) {
|
||||
Particle& p = trap.get_particle(i);
|
||||
arma::vec F = trap.total_force(i, steps * dt);
|
||||
arma::vec a = F / p.mass;
|
||||
p.position += p.velocity * dt;
|
||||
p.velocity += a * dt;
|
||||
}
|
||||
}
|
||||
|
||||
// Initialize Runge-Kutta 4 solver
|
||||
RK4Solver::RK4Solver(PenningTrap& trap_in, double dt_in)
|
||||
: Solver(trap_in, dt_in) {}
|
||||
|
||||
|
||||
// Perform a single RK4 integration step
|
||||
void RK4Solver::step() {
|
||||
int n = trap.num_particles();
|
||||
std::vector<arma::vec> k1_v(n), k1_r(n);
|
||||
std::vector<arma::vec> k2_v(n), k2_r(n);
|
||||
std::vector<arma::vec> k3_v(n), k3_r(n);
|
||||
std::vector<arma::vec> k4_v(n), k4_r(n);
|
||||
|
||||
std::vector<arma::vec> current_positions(n);
|
||||
std::vector<arma::vec> current_velocities(n);
|
||||
|
||||
for (int i = 0; i < n; i++) {
|
||||
current_positions[i] = trap.get_particle(i).position;
|
||||
current_velocities[i] = trap.get_particle(i).velocity;
|
||||
}
|
||||
double t = steps * dt;
|
||||
// Calculate k1
|
||||
for (int i = 0; i < n; i++) {
|
||||
Particle& p = trap.get_particle(i);
|
||||
arma::vec F = trap.total_force(i, t);
|
||||
arma::vec a = F / p.mass;
|
||||
k1_v[i] = a * dt;
|
||||
k1_r[i] = p.velocity * dt;
|
||||
p.position += k1_r[i] / 2;
|
||||
p.velocity += k1_v[i] / 2;
|
||||
}
|
||||
// Calculate k2
|
||||
for (int i = 0; i < n; i++) {
|
||||
Particle& p = trap.get_particle(i);
|
||||
arma::vec F = trap.total_force(i, t + dt / 2);
|
||||
arma::vec a = F / p.mass;
|
||||
k2_v[i] = a * dt;
|
||||
k2_r[i] = p.velocity * dt;
|
||||
p.position = current_positions[i] + k2_r[i] / 2;
|
||||
p.velocity = current_velocities[i] + k2_v[i] / 2;
|
||||
}
|
||||
// Calculate k3
|
||||
for (int i = 0; i < n; i++) {
|
||||
Particle& p = trap.get_particle(i);
|
||||
arma::vec F = trap.total_force(i, t + dt / 2);
|
||||
arma::vec a = F / p.mass;
|
||||
k3_v[i] = a * dt;
|
||||
k3_r[i] = p.velocity * dt;
|
||||
p.position = current_positions[i] + k3_r[i];
|
||||
p.velocity = current_velocities[i] + k3_v[i];
|
||||
}
|
||||
// Calculate k4
|
||||
for (int i = 0; i < n; i++) {
|
||||
Particle& p = trap.get_particle(i);
|
||||
arma::vec F = trap.total_force(i, t + dt);
|
||||
arma::vec a = F / p.mass;
|
||||
k4_v[i] = a * dt;
|
||||
k4_r[i] = p.velocity * dt;
|
||||
p.position = current_positions[i];
|
||||
p.velocity = current_velocities[i];
|
||||
}
|
||||
// Update positions and velocities
|
||||
for (int i = 0; i < n; i++) {
|
||||
Particle& p = trap.get_particle(i);
|
||||
p.position = current_positions[i] + (k1_r[i] + 2.0 * k2_r[i] + 2.0 * k3_r[i] + k4_r[i]) / 6.0;
|
||||
p.velocity = current_velocities[i] + (k1_v[i] + 2.0 * k2_v[i] + 2.0 * k3_v[i] + k4_v[i]) / 6.0;
|
||||
}
|
||||
}
|
||||
|
||||
VelocityVerletSolver::VelocityVerletSolver(PenningTrap& trap_in, double dt_in)
|
||||
: Solver(trap_in, dt_in) {
|
||||
// Initialize accelerations
|
||||
int n = trap.num_particles();
|
||||
accelerations.resize(n);
|
||||
for (int i = 0; i < n; i++) {
|
||||
Particle& p = trap.get_particle(i);
|
||||
arma::vec F = trap.total_force(i, 0.0);
|
||||
accelerations[i] = F / p.mass;
|
||||
}
|
||||
}
|
||||
|
||||
// Perform a single Velocity Verlet integration step
|
||||
void VelocityVerletSolver::step() {
|
||||
int n = trap.num_particles();
|
||||
for (int i = 0; i < n; i++) {
|
||||
Particle& p = trap.get_particle(i);
|
||||
p.position += p.velocity * dt + 0.5 * accelerations[i] * dt * dt;
|
||||
arma::vec new_acceleration = trap.total_force(i, steps * dt) / p.mass;
|
||||
p.velocity += 0.5 * (accelerations[i] + new_acceleration) * dt;
|
||||
accelerations[i] = new_acceleration;
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
// Initialize Analytical solver
|
||||
AnalyticalSolver::AnalyticalSolver(PenningTrap& trap_in, double dt_in)
|
||||
: Solver(trap_in, dt_in) {
|
||||
// Check for validity of analytical solution
|
||||
if (trap_in.get_coulomb_interaction()) {
|
||||
throw std::runtime_error("Analytical solution not valid with Coulomb interactions enabled.");
|
||||
}
|
||||
if (trap_in.get_modified_fields()) {
|
||||
throw std::runtime_error("Analytical solution not valid with modified external fields.");
|
||||
}
|
||||
for (int i = 0; i < trap_in.num_particles(); i++) {
|
||||
if (trap_in.get_particle(i).velocity(0) != 0 || trap_in.get_particle(i).velocity(2) != 0) {
|
||||
throw std::runtime_error("Analytical solution not valid for non-zero initial velocity in x or z direction.");
|
||||
}
|
||||
if (trap_in.get_particle(i).position(1) != 0) {
|
||||
throw std::runtime_error("Analytical solution not valid for non-zero initial position in y direction.");
|
||||
}
|
||||
}
|
||||
|
||||
trap = trap_in;
|
||||
dt = dt_in * constants::second;
|
||||
steps = 0;
|
||||
int n = trap.num_particles();
|
||||
w_0.resize(n);
|
||||
w_z.resize(n);
|
||||
w_plus.resize(n);
|
||||
w_minus.resize(n);
|
||||
A_plus.resize(n);
|
||||
A_minus.resize(n);
|
||||
z_0.resize(n);
|
||||
|
||||
for (int i = 0; i < n; i++) {
|
||||
Particle& p = trap.get_particle(i);
|
||||
double q = p.charge;
|
||||
double m = p.mass;
|
||||
double B0 = trap.B0;
|
||||
double V0 = trap.V0;
|
||||
double d = trap.d;
|
||||
|
||||
w_0[i] = q * B0 / m;
|
||||
w_z[i] = std::sqrt(2 * q * V0 / (m * d * d));
|
||||
double discriminant = w_0[i] * w_0[i] - 2 * w_z[i] * w_z[i];
|
||||
if (discriminant < 0) {
|
||||
throw std::runtime_error("Analytical solution not valid: discriminant is negative.");
|
||||
}
|
||||
w_plus[i] = (w_0[i] + std::sqrt(discriminant)) / 2.0;
|
||||
w_minus[i] = (w_0[i] - std::sqrt(discriminant)) / 2.0;
|
||||
|
||||
arma::vec r = p.position;
|
||||
arma::vec v = p.velocity;
|
||||
|
||||
// Initial conditions
|
||||
double x0 = r(0);
|
||||
double y0 = r(1);
|
||||
double z0 = r(2);
|
||||
double vx0 = v(0);
|
||||
double vy0 = v(1);
|
||||
double vz0 = v(2);
|
||||
|
||||
// Calculate amplitudes
|
||||
A_plus[i] = (vy0 + w_minus[i] * x0) / (w_plus[i] - w_minus[i]);
|
||||
A_minus[i] = -(vy0 + w_plus[i] * x0) / (w_plus[i] - w_minus[i]);
|
||||
z_0[i] = z0;
|
||||
}
|
||||
}
|
||||
|
||||
// Perform a single Analytical "integration" step
|
||||
void AnalyticalSolver::step() {
|
||||
int n = trap.num_particles();
|
||||
double t = steps * dt;
|
||||
for (int i = 0; i < n; i++) {
|
||||
Particle& p = trap.get_particle(i);
|
||||
double x = -A_plus[i] * std::cos(w_plus[i] * t) - A_minus[i] * std::cos(w_minus[i] * t);
|
||||
double y = A_plus[i] * std::sin(w_plus[i] * t) + A_minus[i] * std::sin(w_minus[i] * t);
|
||||
double z = z_0[i] * std::cos(w_z[i] * t);
|
||||
|
||||
double vx = +A_plus[i] * w_plus[i] * std::sin(w_plus[i] * t) + A_minus[i] * w_minus[i] * std::sin(w_minus[i] * t);
|
||||
double vy = A_plus[i] * w_plus[i] * std::cos(w_plus[i] * t) + A_minus[i] * w_minus[i] * std::cos(w_minus[i] * t);
|
||||
double vz = -z_0[i] * w_z[i] * std::sin(w_z[i] * t);
|
||||
|
||||
p.position = arma::vec({x, y, z});
|
||||
p.velocity = arma::vec({vx, vy, vz});
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
// Convert vector of matrices to a 3D cube for easier data handling
|
||||
arma::cube to_cube(const std::vector<arma::mat>& vec) {
|
||||
if (vec.empty()) {
|
||||
return arma::cube();
|
||||
}
|
||||
int rows = vec[0].n_rows;
|
||||
int cols = vec[0].n_cols;
|
||||
int slices = vec.size();
|
||||
arma::cube cube_data(rows, cols, slices);
|
||||
for (size_t i = 0; i < vec.size(); ++i) {
|
||||
cube_data.slice(i) = vec[i];
|
||||
}
|
||||
return cube_data;
|
||||
}
|
||||
|
||||
// Save cube data to a file
|
||||
void save_to_file(const std::string& filename, const arma::cube& data) {
|
||||
data.save(filename, arma::raw_ascii);
|
||||
}
|
||||
|
||||
void save_to_file(const std::string& filename, const arma::vec& data) {
|
||||
arma::mat mat_data = arma::conv_to<arma::mat>::from(data);
|
||||
mat_data.save(filename, arma::raw_ascii);
|
||||
}
|
||||
@@ -0,0 +1,126 @@
|
||||
#include <armadillo>
|
||||
#include "classes.hpp"
|
||||
#include "solvers.hpp"
|
||||
#include "constants.hpp"
|
||||
#include "argparse/argparse.hpp"
|
||||
|
||||
#include <iostream>
|
||||
|
||||
using namespace std;
|
||||
|
||||
int main(int argc, char* argv[]) {
|
||||
argparse::ArgumentParser program("two_particles");
|
||||
program.add_argument("-N", "--steps")
|
||||
.help("Number of time steps")
|
||||
.default_value(4000)
|
||||
.scan<'i', int>();
|
||||
program.add_argument("-n", "--num_particles")
|
||||
.help("Number of particles (either 1 or 2)")
|
||||
.default_value(2)
|
||||
.scan<'i', int>();
|
||||
program.add_argument("-i", "--disable-interactions")
|
||||
.help("Disable Coulomb interactions")
|
||||
.default_value(true)
|
||||
.implicit_value(false);
|
||||
program.add_argument("-E", "--euler")
|
||||
.help("Use Euler solver instead of RK4")
|
||||
.flag();
|
||||
program.add_argument("-A", "--analytical")
|
||||
.help("Use Analytical solver instead of RK4 (only valid for particles without interactions and specific initial conditions)")
|
||||
.flag();
|
||||
program.add_argument("-V", "--velocity-verlet")
|
||||
.help("Use Velocity Verlet solver instead of RK4")
|
||||
.flag();
|
||||
|
||||
try {
|
||||
program.parse_args(argc, argv);
|
||||
} catch (const runtime_error& err) {
|
||||
cerr << err.what() << endl;
|
||||
cerr << program.help().str() << endl;
|
||||
exit(1);
|
||||
}
|
||||
|
||||
double N = program.get<int>("--steps");
|
||||
int n_particles = program.get<int>("--num_particles");
|
||||
bool interactions = program.get<bool>("--disable-interactions");
|
||||
bool use_euler = program.get<bool>("--euler");
|
||||
bool use_analytical = program.get<bool>("--analytical");
|
||||
bool use_velocity_verlet = program.get<bool>("--velocity-verlet");
|
||||
|
||||
double B_0 = 1.0; // Tesla
|
||||
double V_0 = 0.025; // Volt
|
||||
double d = 500e-6; // meter
|
||||
|
||||
PenningTrap trap(B_0, V_0, d);
|
||||
if (interactions) {
|
||||
trap.enable_interactions();
|
||||
}
|
||||
else {
|
||||
trap.disable_interactions();
|
||||
}
|
||||
cout << trap << endl;
|
||||
|
||||
|
||||
Particle p1(arma::vec({20e-6, 0, 20e-6}), arma::vec({0, 25, 0}), 40.078 * constants::amu, 1 * constants::elementary_charge);
|
||||
Particle p2(arma::vec({25e-6, 25e-6, 0}), arma::vec({0, 40, 5}), 40.078 * constants::amu, 1 * constants::elementary_charge);
|
||||
|
||||
trap.add_particle(p1);
|
||||
if (n_particles > 1) {
|
||||
trap.add_particle(p2);
|
||||
}
|
||||
|
||||
double dt = 50e-6 / N; // seconds
|
||||
|
||||
if (use_euler) {
|
||||
EulerSolver solver(trap, dt);
|
||||
solver.simulate(N);
|
||||
solver.save("results/two_particles_" + std::to_string(n_particles) + (interactions ? "_with_interactions_" : "_no_interactions_") + std::to_string(N) + "_steps_euler");
|
||||
vector<arma::mat> positions = solver.get_positions();
|
||||
cout << "Final position of particle 1: " << endl;
|
||||
cout << positions.back().col(0) << endl;
|
||||
if (n_particles > 1) {
|
||||
cout << "Final position of particle 2: " << endl;
|
||||
cout << positions.back().col(1) << endl;
|
||||
}
|
||||
return 0;
|
||||
}
|
||||
|
||||
if (use_analytical) {
|
||||
AnalyticalSolver solver(trap, dt);
|
||||
solver.simulate(N);
|
||||
solver.save("results/two_particles_" + std::to_string(n_particles) + (interactions ? "_with_interactions_" : "_no_interactions_") + std::to_string(N) + "_steps_analytical");
|
||||
vector<arma::mat> positions = solver.get_positions();
|
||||
cout << "Final position of particle 1: " << endl;
|
||||
cout << positions.back().col(0) << endl;
|
||||
if (n_particles > 1) {
|
||||
cout << "Final position of particle 2: " << endl;
|
||||
cout << positions.back().col(1) << endl;
|
||||
}
|
||||
return 0;
|
||||
}
|
||||
|
||||
if (use_velocity_verlet) {
|
||||
VelocityVerletSolver solver(trap, dt);
|
||||
solver.simulate(N);
|
||||
solver.save("results/two_particles_" + std::to_string(n_particles) + (interactions ? "_with_interactions_" : "_no_interactions_") + std::to_string(N) + "_steps_velocity_verlet");
|
||||
vector<arma::mat> positions = solver.get_positions();
|
||||
cout << "Final position of particle 1: " << endl;
|
||||
cout << positions.back().col(0) << endl;
|
||||
if (n_particles > 1) {
|
||||
cout << "Final position of particle 2: " << endl;
|
||||
cout << positions.back().col(1) << endl;
|
||||
}
|
||||
return 0;
|
||||
}
|
||||
|
||||
RK4Solver solver(trap, dt);
|
||||
solver.simulate(N);
|
||||
solver.save("results/two_particles_" + std::to_string(n_particles) + (interactions ? "_with_interactions_" : "_no_interactions_") + std::to_string(N) + "_steps_rk4");
|
||||
vector<arma::mat> positions = solver.get_positions();
|
||||
cout << "Final position of particle 1: " << endl;
|
||||
cout << positions.back().col(0) << endl;
|
||||
if (n_particles > 1) {
|
||||
cout << "Final position of particle 2: " << endl;
|
||||
cout << positions.back().col(1) << endl;
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user